How do I sample N rows from a Pyspark DataFrame?

How do I sample N rows from a PySpark DataFrame? You can use different methods depending on your requirement: sample() for approximate fractions, limit(n) for exact rows, orderBy(rand()).limit(n) for exact random rows, and sampleBy() when you need stratified samples across groups. Each approach serves a different purpose—whether you want speed, precision, or balanced distribution—and choosing the right one depends on your use case.

In this article, we will cover different approaches to sample rows in PySpark while maintaining reproducibility and flexibility.

Key Points-

Multiple Sampling Methods – PySpark provides sample(), limit(n), orderBy(rand()).limit(n), and sampleBy() for flexible sampling.
Approximate vs Exact Sampling – sample() returns an approximate fraction of rows, whereas limit(n) returns exactly N rows.
Random Sampling – orderBy(rand()).limit(n) ensures both randomness and an exact number of rows.
Stratified Sampling – sampleBy() allows sampling while maintaining group-level distributions, useful for imbalanced datasets.
Reproducibility with Seed – All sampling functions support a seed parameter for deterministic results across runs.
Fractional Sampling – sample(fraction=) is ideal for quickly selecting a subset of data without worrying about exact row counts.
Performance Considerations – Sampling large datasets can be optimized by caching the DataFrame or using partitions effectively.
Maintaining Distribution – Using rand() with filters or sampleBy() helps preserve uniform distribution across splits.
Use Cases – Sampling is commonly used for test/train splits, validation sets, exploratory analysis, or reducing dataset size for testing.
Combination of Methods – You can combine methods (e.g., orderBy(rand()) with sampleBy()) for more controlled, reproducible, and balanced sampling.

Using sample() for Fractional Sampling

The sample() function is used to randomly pick a fraction of rows from a DataFrame. It is useful when you want an approximate fraction of rows rather than an exact number.


# Create Spark session
spark = SparkSession.builder.appName("SampleRows").getOrCreate()

# Sample DataFrame
data = [("Sai", 1), ("Ram", 2), ("Geetha", 3), ("David", 4), ("Eva", 5)]
df = spark.createDataFrame(data, ["Name", "Value"])

# Sample 40% of rows randomly
sampled_df = df.sample(withReplacement=False, fraction=0.4, seed=42)
sampled_df.show()

Yields below the output.

withReplacement=False → no duplicate rows.
fraction=0.4 → approximately 40% of rows are selected.
seed=42 → ensures reproducibility across runs.
This method is best when you want a percentage split (e.g., 10%, 50%, etc.).

Use orderBy(rand()) and limit(n) for Random Exact N Rows

To select exactly N random rows, you can combine orderBy(rand()) with limit(n).


from pyspark.sql.functions import rand

# Randomly select 3 rows
random_n_df = df.orderBy(rand()).limit(3)
random_n_df.show()

Yields below the output.

limit(3) ensures exactly 3 rows are returned
rand() generates a random number for each row
orderBy(rand()) shuffles the rows

This is the best method if you need a fixed number of random rows.

Using limit(n) for Exact N Rows

If you want to extract exactly N rows (not random), you can use the limit() function.


# Take exactly 3 rows
exact_n_df = df.limit(3)
exact_n_df.show()

# Output:
# +------+-----+
# |  Name|Value|
# +------+-----+
# |   Sai|    1|
# |   Ram|    2|
# |Geetha|    3|
# +------+-----+

limit(3) always returns exactly 3 rows.
However, the rows are not random → it takes the first N rows from the DataFrame.

Stratified Sampling with sampleBy() (Optional)

If you need to sample rows while preserving distribution across categories, use sampleBy(). This is especially useful for balanced datasets in machine learning tasks.


# Sample dataset
data = [
    (1, "P1"), (1, "P2"), (1, "P3"),
    (2, "P4"), (2, "P5"),
    (3, "P6"), (3, "P7"), (3, "P8"), (3, "P9")
]
df = spark.createDataFrame(data, ["region_id", "property_id"])

# Define fractions per region
fractions = {1: 0.5, 2: 1.0, 3: 0.25}

# Stratified sampling
sampled_df = df.sampleBy("region_id", fractions, seed=42)
sampled_df.show()

# Output:
# +---------+-----------+
# |region_id|property_id|
# +---------+-----------+
# |        2|         P4|
# |        2|         P5|
# |        3|         P8|
# +---------+-----------+

Keeps 50% of region 1, 100% of region 2, and 25% of region 3.
Ensures proportional representation in the sampled dataset.

This is ideal for imbalanced datasets where random sampling may distort class distribution.

Frequently Asked Questions of sample N rows from a Pyspark

You can use the sample() method to select a random fraction of rows.

You can use orderBy(rand()).limit(n) to return exactly N rows chosen randomly.

sample() → returns an approximate fraction of rows randomly.
limit(n) → returns exactly N rows, but always from the top of the DataFrame (not random).

You can use the sampleBy() function to sample rows while preserving the distribution across categories or groups.

Use it when you need exactly N rows randomly instead of an approximate fraction.

Conclusion

In this article, I have covered various methods to sample N rows from a PySpark DataFrame, including sample() for approximate fractions, limit(n) for exact rows, orderBy(rand()).limit(n) for exact random rows, and sampleBy() for stratified sampling across groups. We explored how to maintain reproducibility using seeds, preserve data distribution, and handle imbalanced datasets effectively. By understanding these sampling techniques, you can efficiently create subsets for testing, validation, or exploratory analysis, making PySpark a powerful tool for working with large datasets.

Happy Learning!!