How to get random sample records in PySpark Azure Databricks?

Are you looking to find out how to get a random small dataset from a large dataset of PySpark DataFrame in the Azure Databricks cloud, or maybe you are looking for a sample dataset from a large dataset of PySpark DataFrame Databricks using the sample() and sampleBy() methods? If you are looking for any of these problem solutions, you have landed on the correct page. I will also show you how to use PySpark to get sample datasets in DataFrames in Azure Databricks. I will explain it by taking a practical example. So don’t waste time let’s start with a step-by-step guide to understanding how to get a random sample dataset in PySpark DataFrame.

In this blog, I will teach you the following with practical examples:

  • Syntax of sample(), sampleBy() and takeSample()
  • Random sample records
  • Fixed random sample records
  • Random records with duplication
  • Random sample records based on column
  • Random records using RDD

sample() method is used to produce a random sample dataset of dataframes in PySpark Azure Databricks.

sampleBy() method is used to produce a random sample dataset based on key column of dataframes in PySpark Azure Databricks.

Syntax:

dataframe_name.sample()

dataframe_name.sampleBy()

What is the syntax of the select() function in PySpark Azure Databricks?

The syntax is as follows:

dataframe_name.sample(need_duplication, sample_range, fixed_id)

dataframe_name.sampleBy(need_duplication, sample_range, fixed_id)
Parameter NameRequiredDescription
need_duplication (bool)OptionalIt represents the sample can have duplication or not
sample_range (float)OptionalIt represents the range of sample records
fixed_id (int)OptionalEach fixed_id of sample() function returns a different sample value.
Table 1: sample() and sampleBy() Method in PySpark Databricks Parameter list with Details

Apache Spark Official Documentation Link: sample(), sampleBy() and takeSample()

Create a simple DataFrame

Let’s understand the use of the sample(), sampleBy(), and takeSample() functions with a variety of examples. Let’s start by creating a DataFrame.

Gentle reminder:

In Databricks,

  • sparkSession made available as spark
  • sparkContext made available as sc

In case, you want to create it manually, use the below code.

from pyspark.sql.session import SparkSession

spark = SparkSession.builder 
    .master("local[*]") 
    .appName("azurelib.com") 
    .getOrCreate()

sc = spark.sparkContext

a) Create manual PySpark DataFrame

df=spark.range(100)
df.printSchema()
df.show(5, truncate=False)

"""
root
 |-- id: long (nullable = false)

+---+
|id |
+---+
|0  |
|1  |
|2  |
|3  |
|4  |
+---+
"""

b) Creating a DataFrame by reading files

Download and use the below source file.

# replace the file_path with the source file location which you have downloaded.

df_2 = spark.read.format("csv").option("inferSchema", True).option("header", True).load(file_path)
df_2.printSchema()

"""
root
 |-- id: string (nullable = true)
"""

Note: Here, I will be using the manually created DataFrame.

How to get sample records in PySpark Azure Databricks?

The sample() function in PySpark produces a range of random sample records from large datasets by providing the fraction value. This is not guaranteed to provide exactly the fraction specified of the total count of the given DataFrame.

Example 1:

# First run
print(f"First run count: {df.sample(fraction=0.1).count()}")
sample_df.show()
"""
Output:

First run count: 9
+---+
| id|
+---+
| 17|
| 27|
| 51|
| 64|
| 71|
| 73|
| 77|
| 81|
| 83|
| 84|
| 99|
+---+

"""

Example 2:

# Second run
print(f"Second run count: {df.sample(fraction=0.1).count()}")
sample_df.show()

"""
Output:

Second run count: 8
+---+
| id|
+---+
| 17|
| 27|
| 51|
| 64|
| 71|
| 73|
| 77|
| 81|
| 83|
| 84|
| 99|
+---+

"""

How to get same sample records in PySpark Azure Databricks?

The sample() function in PySpark produces a range of random sample records from large datasets by providing the fraction value for each run. But what in case you need a fixed sample record for comparison? This can be achieved by specifying the seed parameter to a fixed value.

Example 1:

sample_df = df.sample(fraction=0.1, seed=324)
# Note: seed value, ensure the giving the same sample data on each run

# First run
print(f"First run count: {sample_df.count()}")
sample_df.show()

"""
Output:

First run count: 5
+---+
| id|
+---+
| 21|
| 22|
| 31|
| 72|
| 88|
+---+

"""

Example 2:

# Second run
print(f"Second run count: {sample_df.count()}")
sample_df.show()

"""
Output:

Second run count: 5
+---+
| id|
+---+
| 21|
| 22|
| 31|
| 72|
| 88|
+---+

"""

Example 3:

sample_df_2 = df.sample(fraction=0.1, seed=456)
print(f"Third run count with another seed value: {sample_df_2.count()}")
sample_df_2.show()

"""
Output:

Third run count with another seed value: 11
+---+
| id|
+---+
| 19|
| 21|
| 35|
| 39|
| 48|
| 58|
| 60|
| 66|
| 72|
| 79|
| 95|
+---+

"""

How to get duplicate sample records in PySpark Azure Databricks?

The sample() function in PySpark produces a range of random sample records from large datasets by providing the fraction value for each run. But what in case you need a duplicate record for comparison? This can be achieved by specifying the withReplacement parameter to a fixed value.

Example:

sample_df_3 = df.sample(withReplacement=True, fraction=0.2, seed=124)
print(f"First duplicate run output: {', '.join([str(each.id) for each in sample_df_3.collect()])}")

sample_df_4 = df.sample(withReplacement=True, fraction=0.2, seed=133)
print(f"Second duplicate run output: {', '.join([str(each.id) for each in sample_df_4.collect()])}")

"""
Output:

First duplicate run output: 0, 2, 2, 3, 5, 6, 7, 22, 25, 26, 27, 32, 35, 39, 42, 45, 56, 56, 62, 80, 83, 86, 94, 95, 97
Second duplicate run output: 14, 15, 20, 36, 47, 51, 52, 59, 60, 62, 70, 75, 76, 77, 81, 81, 85, 96, 98

"""

As you can that, both the first(2) and second(81) output has duplicate values.

How to get sample records based on columns in PySpark Azure Databricks?

The sampleBy() function in PySpark produces a range of random sample records from large datasets by key column.

f2 = df.withColumn("key", df.id % 3).select("key", "id")
print("Grouping the dataframe based on 'key' column:")
df2.groupBy("key").count().show()

# 10% of 34 records = 3.4
df3 = df2.sampleBy("key", fractions={0: 0.1, 1: 1}, seed=123)
print("Number of sample records reproduced based on fractions:")
df3.groupBy("key").count().show()

"""
Output:

Grouping the dataframe based on 'key' column:
+---+-----+
|key|count|
+---+-----+
|  0|   34|
|  1|   33|
|  2|   33|
+---+-----+

Number of sample records reproduced based on fractions:
+---+-----+
|key|count|
+---+-----+
|  1|   33|
|  0|    4|
+---+-----+

"""

How to get sample records from RDD in PySpark Azure Databricks?

You can use sample() and takeSample() functions in PySpark RDD. Whereas sample() is a transformation and takeSample() is an action, so use takeSample() carefully as it sends the results to the driver memory may lead to Out of Memory issue.

rdd = spark.sparkContext.range(0,100)

# 1. sample()
rdd.sample(withReplacement=False, fraction=0.1,seed=1).collect()

# 2. takeSample() -> action -> send result to driver
rdd.takeSample(withReplacement=False, num=10,seed=1)

"""
Output:

[11, 13, 49, 61, 65]
[93, 22, 85, 39, 92, 96, 49, 61, 6, 29]

"""

I have attached the complete code used in this blog in a notebook format to this GitHub link. You can download and import this notebook in databricks, jupyter notebook, etc.

When should you go for sampling in PySpark Azure Databricks?

These could be the possible reasons:

  1. To get sample records
  2. To get random sample records
  3. To get random sample records with duplication
  4. To get random sample records without duplication
  5. To get sample records based on defined column

Real World Use Case Scenarios for sampling in PySpark Azure Databricks?

In general, a data scientist or data analyst deals with huge datasets that lead to time-consuming data processing. You can take a sample of the huge dataset for data analysis, in this case, you can go for data sampling.

What are the alternatives to the sample() function in PySpark Azure Databricks?

These alternatives were discussed with multiple examples in the above section.

  • sample() used for getting random sample records
  • sampleBy() is used for getting random sample records based on columns
  • takeSample() used for getting random sample records from RDD

Final Thoughts

In this article, we have learned about the PySpark sample(), sampleBy(), and takeSample() methods to select the columns of a DataFrame in Azure Databricks along with the examples explained clearly. I have also covered different scenarios with practical examples that could be possible. I hope the information that was provided helped in gaining knowledge.

Please share your comments and suggestions in the comment section below and I will try to answer all your queries as time permits.

PySpark in Azure Databricks, as explained by Arud Seka Berne S on azurelib.com.

As a big data engineer, I design and build scalable data processing systems and integrate them with various data sources and databases. I have a strong background in Python and am proficient in big data technologies such as Hadoop, Hive, Spark, Databricks, and Azure. My interest lies in working with large datasets and deriving actionable insights to support informed business decisions.