Explain sampleBy() function in PySpark

In PySpark, the sampleBy() The function is used to perform stratified sampling based on the values of a specific column in a DataFrame. Unlike simple random sampling, stratified sampling allows you to control the fraction of data selected from each group.

For example, if your dataset contains customers grouped by country, you might want to select 10% from the U.S., 20% from India, and 5% from other countries.

In simple terms, stratified sampling divides the dataset into strata (groups) based on given criteria, and then extracts samples from each stratum according to the specified fraction.

In this article, I will explain how the sampleBy() function generates stratified samples without replacement.

Key Points-

Used for stratified sampling in PySpark.
Works on one column at a time; for multiple columns, create a composite key.
Requires a fractions dictionary to define the sampling percentage for each group.
If a stratum is not specified, it takes zero as the default.
Setting the fraction to 1.0 selects all rows from that group.
The original DataFrame remains unchanged; it returns a new sampled DataFrame.
Supports an optional seed parameter for reproducible results.

PySpark sampleBy() Function

The sampleBy() function in PySpark is used to take samples from different groups in your DataFrame. You choose a column to define the groups and decide what fraction of rows to pick from each group.

Syntax

Below is the syntax of the sampleBy() function.


# Syntax of the sampleBy() function
DataFrameStatFunctions.sampleBy(col, fractions, seed=None)

Parameters

col: – column name from DataFrame
fractions: – It’s Dictionary type takes key and value, Example: {"A": 0.1, "B": 0.2}
seed: - Seed for sampling (default a random seed). Used to reproduce the same random sampling.

Return Value

Returns a new DataFrame with stratified samples based on the given fractions, leaving the original DataFrame unchanged.

Usage of PySpark sampleBy

You can use the sampleBy() function to get stratified samples of each group based on the given fractions. First, create a DataFrame with a new column that labels each number as even or odd. Next, apply stratified sampling with sampleBy() using the defined fractions. Finally, use groupBy() and count() to display the sampled rows for each group.

In this example, we’ll divide numbers into even and odd groups, and then apply stratified sampling to pick 10% of even numbers and 20% of odd numbers.


# Get stratified sample for 10% even and 20% odd
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when

# Create Spark session
spark = SparkSession.builder \
    .master("local[1]") \
    .appName("AIStudyLabs.com") \
    .getOrCreate()

# Create DataFrame with even/odd groups
dataset = spark.range(0, 100).withColumn(
    "group",
    when(col("id") % 2 == 0, "even").otherwise("odd")
)
dataset.show()

# Stratified sampling: 10% even, 20% odd
fractions = {"even": 0.1, "odd": 0.2}

sampled = dataset.sampleBy("group", fractions, seed = 1)
sampled.groupBy("group").count().orderBy("group").show()

Yields below the output.

Select All Rows from One Group

You can also get the total number of rows from a specific group. To do so, you can set the fraction for that group to 1.0 and pass it into this function to obtain the count of all rows in the specified group using the count() function.


# Take 100% of rows from the "even" group
dataset.sampleBy("group", fractions={"even": 1.0}, seed=0).count()

# Output:
# 50

Since half of the numbers between 0–99 are even.

Stratified Sampling with Multiple Columns

So far, we have seen how to use PySpark’s sampleBy() function on a single column. Now, let’s explore how to apply it when working with multiple columns.

To perform this, you can combine the values of multiple columns into a single composite key. This composite key represents the unique combination of categories across those columns.

Once the composite key is created, you can define different sampling fractions for each combined group and pass them to the sampleBy() function. Spark will then apply stratified sampling according to the specified fractions. After sampling, you can use operations like groupBy() and count() to verify how many rows were selected for each composite group.


# Get stratified samples with multiple columns
from pyspark.sql.functions import concat_ws

# Add another grouping column
dataset = dataset.withColumn("mod3", col("id") % 3)

# Create composite key
dataset = dataset.withColumn("group_mod3", concat_ws("_", col("group"), col("mod3")))

# Sample different fractions for composite groups
fractions = {"even_0": 0.2, "odd_1": 0.3}
sampled = dataset.sampleBy("group_mod3", fractions, seed=42)
sampled.groupBy("group_mod3").count().orderBy("group_mod3").show()

Yields below output.


# Output:
+----------+-----+
|group_mod3|count|
+----------+-----+
|    even_0|    6|
|     odd_1|    4|
+----------+-----+

If a Fraction Is Not Specified for a Group

If you don’t specify a stratum in the fractions dictionary, Spark treats it as 0, which means no rows will be sampled from that group.


# Get the stratified sample when on of stratum is 0
from pyspark.sql.functions import when, col

# Create DataFrame with even/odd groups
dataset = spark.range(0, 10).withColumn(
    "group",
    when(col("id") % 2 == 0, "even").otherwise("odd")
)

# Fractions dictionary only includes "even"
fractions = {"even": 0.5}

# Apply stratified sampling
sampled = dataset.sampleBy("group", fractions, seed=1)

# Count rows per group after sampling
sampled.groupBy("group").count().orderBy("group").show()

Yields below output.


# Output:
+-----+-----+
|group|count|
+-----+-----+
| even|    2|
+-----+-----+

Key Differences Between sample() and sampleBy()

The table below provides a complete comparison of the features ofsample() and sampleBy().

| Feature | sample() | sampleBy() | | --- | --- | --- | | Type of Sampling | Simple random sampling | Stratified sampling | | Groups | Ignores groups | Considers groups (strata) | | Fraction | Applies to the entire DataFrame | Fraction can vary per group | | Replacement | Can be with or without replacement | Always without replacement | | Use Case | Random test/train splits, quick sampling | Balanced group-wise sampling |

Frequently Asked Questions of PySpark sampleBy()

sample() → Performs simple random sampling on the entire DataFrame.
sampleBy() → Performs stratified sampling, letting you control how many rows to sample from each group (stratum).

By default, sampleBy() samples without replacement.
If you need sampling with replacement, you should use the sample() function instead.

Any group not included in the fractions dictionary is treated as 0, meaning no rows are sampled from that group.

sampleBy() works with one column at a time.
For multiple columns, you can create a composite key (e.g., by concatenating column values) and then apply sampleBy().

The seed parameter ensures reproducibility.
Without it, sampling results will differ every time you run the code.

Conclusion

In this article, I have explained how to use the sampleBy() function in PySpark for stratified sampling. Unlike simple random sampling, sampleBy() gives you full control over how many rows to pick from each group.

We saw how to create even/odd groups and select different fractions from each.
We learned how to select all rows from one group by setting the fraction to 1.0.
Finally, we covered how to handle multiple columns by creating a composite key for stratified sampling.

Use sampleBy() when you need balanced, controlled sampling across groups in your dataset, such as customers by region, transactions by category, or test/train splits with class proportions preserved.

Happy Learning!!

Reference

https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrameStatFunctions.sampleBy.html