In PySpark, the randomSplit() function is used to divide a DataFrame into multiple smaller DataFrames based on specified weights. Unlike simple row selection, randomSplit() splits the dataset into training, validation, and test sets or any other proportion you need.
For example, if your dataset contains 1 million rows, you might want to split it into 70% training data, 20% validation data, and 10% test data.
In simple terms, randomSplit() randomly distributes rows across multiple DataFrames according to the fractions you provide.
In this article, we will learn how the randomSplit() function works in PySpark and how you can use it to split DataFrames efficiently into training, validation, and test datasets.
Key Points-
- Used for randomly splitting a DataFrame into multiple smaller DataFrames.
- Works based on weights provided for each output DataFrame.
- Returns a list of DataFrames, leaving the original DataFrame unchanged.
- Supports an optional
seedparameter for reproducible splits. - It can be used for splitting datasets into training, validation, and test sets.
- If your dataset is imbalanced, always use
sampleBy()(stratified sampling) to preserve the ratio of classes across splits.
<a id="h-pyspark-randomsplit-function"></a>
PySpark randomSplit() Function
The randomSplit() function splits a DataFrame into multiple smaller DataFrames according to the weights you provide.
<a id="h-syntax"></a>
Syntax
Below is the syntax of the randomSplit() function.
# Syntax of randomSplit() function
DataFrame.randomSplit(weights, seed=None)
<a id="h-parameters"></a>
Parameters
weights:List of numbers representing the relative sizes of the splits. Example:[0.7, 0.2, 0.1].seed:Optional seed for reproducibility. Without it, splits will differ each time you run the code.
Return Value
Returns a list of DataFrames split according to the specified weights.
Usage of PySpark randomSplit()
You can use randomSplit() to split a DataFrame into multiple subsets. By providing a list of weights corresponding to training, validation, and test sets, Spark randomly returns rows to each split according to the specified ratios.
For example, to split a dataset into 70% training, 20% validation, and 10% test set.
# Use randomsplit() to divide DattaFrame into multiple parts
from pyspark.sql import SparkSession
# Create Spark session
spark = SparkSession.builder \
.master("local[1]") \
.appName("AIStudyLabs.com") \
.getOrCreate()
# Create sample DataFrame
data = spark.range(0, 100)
data.show()
# Split the DataFrame: 70% train, 20% validation, 10% test
train_df, val_df, test_df = data.randomSplit([0.7, 0.2, 0.1], seed=42)
# Show the counts of each split
print("Train count:", train_df.count())
print("Validation count:", val_df.count())
print("Test count:", test_df.count())
Yields below the output.
The above code creates a DataFrame with 100 rows and splits it into three DataFrames according to the given ratios. Note that counts may not match exactly because the split is random; fractions are approximate.
<a id="h-using-randomsplit-for-training-and-test-sets"></a>
Using randomSplit() for Training and Test Sets
The most common use case of randomSplit() is dividing a dataset into training and test sets for machine learning. To achieve this, you pass two weights into the function along with a seed value, which gives you the training and test DataFrames.
# Split DataFrame into training (80%) and test (20%)
train_df, test_df = data.randomSplit([0.8, 0.2], seed=1)
# Show first 5 rows of training set
train_df.show(5)
Yields below the output.
Here, the dataset is divided into 80% training and 20% testing. The seed ensures reproducibility, so you get the same split every time you run the code.
Using randomSplit() with Multiple DataFrames
You can also divide a DataFrame into more than two or three parts if needed. By providing a list of weights to the function, Spark will randomly create multiple split DataFrames based on those proportions.
# Split into four DataFrames
df1, df2, df3, df4 = data.randomSplit([0.4, 0.3, 0.2, 0.1], seed=10)
print("DF1 count:", df1.count())
print("DF2 count:", df2.count())
print("DF3 count:", df3.count())
print("DF4 count:", df4.count())
# Output:
# DF1 count: 33
# DF2 count: 32
# DF3 count: 24
# DF4 count: 11
In this example, the dataset is divided into 40%, 30%, 20%, and 10%. PySpark normalizes weights automatically, so they don’t need to sum to 1.
<a id="h-handling-imbalanced-datasets"></a>
Handling Imbalanced Datasets
When working with imbalanced datasets, randomSplit() can produce unexpected results because it doesn’t preserve class distribution.
from pyspark.sql import functions as F
# Create imbalanced dataset
data = [(i, 0) for i in range(991)] + [(i, 1) for i in range(991, 1000)]
df = spark.createDataFrame(data, ["id", "is_clicked"])
# Split dataset
train_df, test_df = df.randomSplit([0.7, 0.3], seed=42)
print("Original counts:", df.groupBy("is_clicked").count().collect())
print("Train counts:", train_df.groupBy("is_clicked").count().collect())
print("Test counts:", test_df.groupBy("is_clicked").count().collect())
# Output:
# Original counts: [Row(is_clicked=0, count=991), Row(is_clicked=1, count=9)]
# Train counts: [Row(is_clicked=0, count=738), Row(is_clicked=1, count=6)]
# Test counts: [Row(is_clicked=0, count=253), Row(is_clicked=1, count=3)]
You may notice that after splitting, the is_clicked column (imbalanced target variable) doesn’t appear evenly distributed across the train and test sets. This happens because randomSplit() doesn’t preserve class distribution; it only randomly distributes rows.
<a id="h-fix-distribution-issues"></a>
Fix Distribution Issues
To preserve distribution across splits, you can try the following approaches:
<a id="h-add-a-random-column-and-sort-before-splitting"></a>
Add a Random Column and Sort Before Splitting
You can use the rand() function to generate a random column, sort the DataFrame by it, and then apply randomSplit() to achieve a more uniform distribution.
# Sample DataFrame
data = [(i, i % 2) for i in range(1, 11)] # (id, is_clicked) with 0/1
df = spark.createDataFrame(data, ["id", "is_clicked"])
# Add random column and split
df = df.withColumn("rand", F.rand(seed=42)).orderBy("rand")
train_df, test_df = df.randomSplit([0.7, 0.3], seed=42)
print("Training Data:")
train_df.show()
print("Testing Data:")
test_df.show()
# Output
# Training Data:
# +---+----------+-------------------+
# | id|is_clicked| rand|
# +---+----------+-------------------+
# | 1| 1| 0.619189370225301|
# | 2| 0| 0.5096018842446481|
# | 4| 0|0.26322809041172357|
# | 5| 1| 0.6702867696264135|
# | 6| 0| 0.5173283545794627|
# | 8| 0|0.06993233728279169|
# +---+----------+-------------------+
# Testing Data:
# +---+----------+------------------+
# | id|is_clicked| rand|
# +---+----------+------------------+
# | 3| 1|0.8325259388871524|
# | 7| 1|0.9991441647585968|
# | 9| 1|0.9696695610826327|
# | 10| 0|0.7959575617927873|
# +---+----------+------------------+
<a id="h-use-filter-instead-of-randomsplit"></a>
Use Filter Instead of randomSplit
You can also split your dataset into train and test sets by creating a random column with rand() and applying filter conditions. This approach ensures that the distribution remains uniform without relying on randomSplit().
# Sample DataFrame
data = [(i, i % 2) for i in range(1, 11)] # (id, is_clicked) with 0/1
df = spark.createDataFrame(data, ["id", "is_clicked"])
# Add random column
df = df.withColumn("rand", F.rand(seed=42))
# Split using filter
train_df = df.filter(df.rand < 0.7)
test_df = df.filter(df.rand >= 0.7)
print("Training Data:")
train_df.show()
print("Testing Data:")
test_df.show()
# Output:
# Training Data:
# +---+----------+-------------------+
# | id|is_clicked| rand|
# +---+----------+-------------------+
# | 1| 1| 0.619189370225301|
# | 2| 0| 0.5096018842446481|
# | 4| 0|0.26322809041172357|
# | 5| 1| 0.6702867696264135|
# | 6| 0| 0.5173283545794627|
# | 8| 0|0.06993233728279169|
# +---+----------+-------------------+
# Testing Data:
# +---+----------+------------------+
# | id|is_clicked| rand|
# +---+----------+------------------+
# | 3| 1|0.8325259388871524|
# | 7| 1|0.9991441647585968|
# | 9| 1|0.9696695610826327|
# | 10| 0|0.7959575617927873|
This gives you a cleaner and reproducible split with uniform distribution.
<a id="h-use-sampleby-for-stratified-splits"></a>
Use sampleBy() for Stratified Splits
You can use the sampleBy() function to preserve the class distribution when splitting imbalanced datasets. Here, we keep 70% of each class (0 and 1) in the training set, and the rest goes into the test set.
# Sample DataFrame
data = [(i, i % 2) for i in range(1, 11)] # (id, is_clicked) with 0/1
df = spark.createDataFrame(data, ["id", "is_clicked"])
# Stratified sampling: Keep 70% of each class in training set
fractions = {0: 0.7, 1: 0.7}
train_df = df.sampleBy("is_clicked", fractions, seed=42)
test_df = df.subtract(train_df)
print("Training Data:")
train_df.show()
print("Testing Data:")
test_df.show()
# Output:
# Training Data:
# +---+----------+
# | id|is_clicked|
# +---+----------+
# | 1| 1|
# | 2| 0|
# | 4| 0|
# | 5| 1|
# | 6| 0|
# | 8| 0|
# +---+----------+
# Testing Data:
# +---+----------+
# | id|is_clicked|
# +---+----------+
# | 7| 1|
# | 9| 1|
# | 3| 1|
# | 10| 0|
# +---+----------+
This ensures both splits maintain the same class distribution.
Key Differences Between sample() and randomSplit()
The table below provides a complete comparison of the features of sample() and randomSplit().
| Feature | sample() | randomSplit() | | --- | --- | --- | | Type of Operation | Random sampling | Random splitting | | Output | Single DataFrame | Multiple DataFrames (list) | | Fraction | Single fraction | List of weights | | Replacement | Can be with/without replacement | Always without replacement | | Use Case | Quick random sampling | Train/test/validation splits |
<a id="h-frequently-asked-questions-of-pyspark-randomsplit"></a>
Frequently Asked Questions of PySpark randomSplit()
It divides a DataFrame into multiple smaller DataFrames based on the weights you provide. The split is random and approximate, not exact.
Nothing special — Spark automatically normalizes them. For example, [7, 2, 1] is treated the same as [0.7, 0.2, 0.1].
randomSplit() → splits the DataFrame into parts randomly, without considering group distribution.
sampleBy() → performs stratified sampling, ensuring each group is sampled proportionally to the given fractions.
randomSplit() does not preserve class distribution. For imbalanced datasets, use sampleBy() to perform stratified sampling and maintain balance across splits.
The underlying partitions in a DataFrame do not guarantee order. This can sometimes lead to overlapping or clustered splits. To fix it, add a random column with rand() and sort the DataFrame before splitting, or use filter() instead of randomSplit().
<a id="h-conclusion"></a>
Conclusion
In this article, we learned how to use PySpark’s randomSplit() function to split DataFrames randomly into training, validation, and test sets.
We also explored:
- How to use multiple splits.
- Why imbalanced datasets may cause issues.
- How to fix anomalies using sorting, filtering, or caching.
- Differences between
randomSplit()andsampleBy().
Use randomSplit() for quick, probabilistic splits, and sampleBy() when class distribution needs to be preserved.
Happy Learning!!
<a id="h-reference"></a>
Reference
<a id="h-related-articles"></a>