RDD - Word Count Example - Introduction to PySpark for Data Engineering

RDD - Word Count Example

Word count example is a most asked interview question. In this lesson, we'll count the occurrences of each word in a text file. Here are the steps:

First, create a text file containing the text you want to analyze. For this example, let's assume you have a text file named test.txt with the following content:

Project Gutenberg’s
Alice’s Adventures in Wonderland
Project Gutenberg’s
Adventures in Wonderland
Project Gutenberg’s

Now, let's create a PySpark session and read this file into an RDD.

# Imports
from pyspark.sql import SparkSession

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

# Load the input text file into an RDD
rdd = spark.sparkContext.textFile("test.txt")

Now, let's follow the below steps to get the unique words and count of each word.

1. flatMap() Transformation

flatMap() transformation flattens the RDD after applying the function and returns a new RDD. In the below example, first, it splits each record by space in an RDD and finally flattens it. The resulting RDD consists of a single word on each record.

# Split each line into words and flatten the result
rdd2 = rdd.flatMap(lambda x: x.split(" "))

2. map() Transformation

map() transformation is used to apply any complex operations like adding a column, or updating a column etc., the output of map transformations would always have the same number of records as the input.

In our word count example, we are adding a new column with a value of 1 for each word, the result of the RDD contains key-value pairs, with the word of type String as Key, and 1 of type Int as value.

# Map each word to a key-value pair (word, 1)
rdd3 = rdd2.map(lambda x: (x, 1))

3. reduceByKey() Transformation

reduceByKey() merges the values for each key with the function specified. In our example, it reduces the word string by applying the sum function to the value. The result of our RDD contains unique words and their count.

# Reduce by key to count the occurrences of each word
rdd4 = rdd3.reduceByKey(lambda x, y: x + y)

4. Collect and Print the Result

Finally, collect the results and print them to the console.

# Collect the result
result = rdd4.collect()

# Print the result to console
for word, count in result:
    print(f"{word} -> {count}")

Alternatively, you can write this RDD into a text file using the saveAsTextFile() method.

# Write final RDD back to text file
rdd4.saveAsTextFile("/tmp/word-count.txt")

In this example, we used PySpark to read the input text file, split the lines into words, and then performed a word count using RDD transformations like flatMap, map, and reduceByKey. Finally, we collected and printed the results.

Keep in mind that this is a basic example. In a real-world scenario, you might want to clean and preprocess the data further, handle punctuation, and perform more advanced text analysis tasks as needed.

Complete Word Count Example

Following is a complete example of a word count by using several RDD transformations.

# Imports
from pyspark.sql import SparkSession

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

# Load the input text file into an RDD
rdd = spark.sparkContext.textFile("test.txt")

# Split each line into words and flatten the result
rdd2 = rdd.flatMap(lambda x: x.split(" "))

# Map each word to a key-value pair (word, 1)
rdd3 = rdd2.map(lambda x: (x, 1))

# Reduce by key to count the occurrences of each word
rdd4 = rdd3.reduceByKey(lambda x, y: x + y)

# Collect the result
result = rdd4.collect()

# Print the result to console
for word, count in result:
    print(f"{word} -> {count}")

This example yields the below output:

Gutenberg’s -> 3 
Wonderland -> 2   
in -> 2 
Alice’s -> 1 
Project -> 3 
Adventures -> 2

RDD - Word Count Example