ArrowBack to blog

Getting Started with PySpark, Kafka, SQL, and AI

How to run SQL queries over Kafka in seconds? With PySpark, we get a system where we can run ad-hoc analysis locally and any SQL queries to explore our data quickly.

Author's avatar
Stéphane Derosiaux
September 17th, 2023
Blog's image

PySpark is a Python interface to write Apache Spark applications to use it in command line. The Apache Spark platform is built to crunch big datasets in a distributed way.

Combining PySpark with Kafka, we get a system that can process real-time data from Kafka in seconds using PySpark commands. Running ad-hoc analysis locally and various SQL queries to explore our data quickly is particularly useful.

Let's see how to:

  • Get started with PySpark and Kafka
  • Do filtering and transformations using SQL in pyspark
  • Do filtering and SQL transformations with Conduktor
  • Finish with some madness using OpenAI/ChatGPT to generate random JSON into Kafka!

Fasten your seatbelt, we're going to explore a lot of things!

Install PySpark# 

If you're running on macOS, use Homebrew to install PySpark. Check Apache Spark website for other systems or installation methods.

1$ brew install apache-spark

This will install a ton of things! Spark is a big project, with many components and relies on the JVM and the Scala programming language at its core. We're only interested in the pyspark component here, which is the Python CLI for Spark.

pyspark should be available in your terminal (open a new one otherwise):

1$ pyspark
2Python 3.11.5 (main, Aug 24 2023, 15:09:45) [Clang 14.0.3 (clang-1403.0.22.14.1)] on darwin
3Type "help", "copyright", "credits" or "license" for more information.
4Setting default log level to "WARN".
5To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).
623/09/14 00:13:14 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable
7Welcome to
8      ____              __
9     / __/__  ___ _____/ /__
10    _\ \/ _ \/ _ `/ __/  '_/
11   /__ / .__/\_,_/_/ /_/\_\   version 3.4.1
12      /_/
13
14Using Python version 3.11.5 (main, Aug 24 2023 15:09:45)
15Spark context Web UI available at http://stephane.lan:4040
16Spark context available as 'sc' (master = local[*], app id = local-1694643195574).
17SparkSession available as 'spark'.
18>>> ▓

Wonderful, it works! We have a local Spark instance running, powered by Python, with a local Spark UI available at http://localhost:4040. We'll do some fancy things with this UI later in this article, feel free to explore.

Spark UI

Let's familiarize ourselves with the pyspark API first by doing some basic operations, like creating a dataset and printing its schema:

1>>> data = [("Java", "20000"), ("Python", "100000"), ("Scala", "3000")]
2>>> df = spark.createDataFrame(data)
3df.show()
4>>> df.show()
5+------+------+
6|    _1|    _2|
7+------+------+
8|  Java| 20000|
9|Python|100000|
10| Scala|  3000|
11+------+------+
12>>> df.printSchema()
13root
14 |-- _1: string (nullable = true)
15 |-- _2: string (nullable = true)

A DataFrame (df for short) is a Spark concept. It's like a table in a spreadsheet. It has rows and columns, and each column has a name. Each row is a record that contains data for each column. For example, in a table storing information about books, a row could contain data about a single book: its title, author, and price.

In our test above, we created a DataFrame with 3 rows and 2 unnamed columns.

Just like you can sort, filter, or perform calculations on a table in a spreadsheet program, you can do similar operations on a DataFrame in a programming environment. The main difference is that DataFrames are designed to handle large (millions, billions) volumes of data efficiently, even data that won't fit into your computer's memory.

That's exactly why Kafka and Spark are a great match! 🚀

Using Spark to read data from Kafka# 

We can connect our pyspark to Kafka to read data.

First, we need a Kafka! If we don't have one already running, we can use Upstash to have a free Kafka cluster in the cloud. Conduktor will be used to create topics and produce data, and pyspark to do some transformations.

Follow this guide to get your free Kafka cluster up and running in a few minutes: Getting Started with Conduktor and Upstash.

If you don't have installed Conduktor yet, install it and add the configuration of your Upstash cluster.

1$ curl -L https://releases.conduktor.io/console -o docker-compose.yml && docker compose up

Conduktor

Then we create a topic "hello" and produce some data into it using our Kafka UI:

Conduktor

All good from a Kafka and Conduktor perspective. Let's go back to Spark.

We start pyspark with the Kafka dependency to be able to read from Kafka:

1# If we don't add the spark-sql-kafka package, we will run into this error:
2# pyspark.errors.exceptions.captured.AnalysisException: Failed to find data source: kafka.
3
4$ pyspark --packages org.apache.spark:spark-sql-kafka-0-10_2.12:3.4.1

Then, we use pyspark to read from our topic "hello" for 10 seconds. In the code below, replace the XXX with our own credentials from Upstash for the kafka.sasl.jaas.config property. ⚠️

Notice the important "startingOffsets": "earliest" option: it tells Spark to start from the beginning of the topic, and not from the end (which is the default behavior).

Copy this whole block into pyspark (don't forget to update with your own credentials):

1from pyspark.sql import SparkSession
2from datetime import datetime, timedelta
3import time
4
5# We will keep the consumer running for 10 seconds
6stop_time = datetime.now() + timedelta(seconds=10)  # 10 seconds from now
7
8##
9## TODO TO UPDATE kafka.sasl.jaas.config
10##
11kafka_options = {
12    "kafka.bootstrap.servers": "romantic-drake-10214-eu1-kafka.upstash.io:9092",
13    "kafka.sasl.mechanism": "SCRAM-SHA-256",
14    "kafka.security.protocol": "SASL_SSL",
15    "kafka.sasl.jaas.config": """org.apache.kafka.common.security.scram.ScramLoginModule required username="XXX" password="YYY";""",
16    "startingOffsets": "earliest", # Start from the beginning when we consume from kafka
17    "subscribe": "hello"           # Our topic name
18}
19
20# Subscribe to Kafka topic "hello"
21df = spark.readStream.format("kafka").options(**kafka_options).load()
22
23# Deserialize the value from Kafka as a String for now
24deserialized_df = df.selectExpr("CAST(value AS STRING)")
25
26# Query Kafka and wait 10sec before stopping pyspark
27query = deserialized_df.writeStream.outputMode("append").format("console").start()
28time.sleep(10)
29query.stop()

We should see the following output in pyspark, while we produce data in the topic hello using Conduktor:

1+---------------+
2|          value|
3+---------------+
4|      my       |
5|      first    |
6|      message  |
7...
8+---------------+
9only showing top 20 rows
10>>>

The connection works! Now, let's explore transformations (mappings, filters, etc.) using JSON (structured data, easier to work with) using pyspark and Kafka.

Tips: Spark UI to see active queries# 

As working in a shell might be painful (we start streaming queries reading from Kafka and forget to .stop() them) to know what is running and what is not, we can use the Spark UI anytime to see the active queries and their status:

Example http://localhost:4040/StreamingQuery/:

Spark UI

Transforming JSON from Kafka with SQL# 

We're going the Netflix way and will work with JSON representing "view events", like:

1{
2  "user_id": 18,
3  "content": "Stranger Things",
4  "watched_at": "2023-08-10 10:00:01",
5  "rating": 5
6}

We will produce JSON data (the view events) into Kafka using Conduktor, and read it using pyspark. At the same time, we'll transform the incoming data SQL as it's the ubiquitous language to work with structured data.

Produce the JSON above in our topic `hello' then execute the following code:

1from pyspark.sql.functions import *
2from pyspark.sql.types import *
3
4json_schema = StructType([
5  StructField("user_id", StringType()),
6  StructField("content", StringType()),
7  StructField("watched_at", TimestampType()),
8  StructField("rating", IntegerType()),
9])
10
11# df is our DataFrame reading from Kafka above
12json_df = df.select(from_json(col("value").cast("string"), json_schema).alias("value"))
13json_df.printSchema()
14# root
15#  |-- value: struct (nullable = true)
16#  |    |-- user_id: string (nullable = true)
17#  |    |-- content: string (nullable = true)
18#  |    |-- watched_at: timestamp (nullable = true)
19#  |    |-- rating: integer (nullable = true)
20
21query = json_df.writeStream.outputMode("append").format("console").start()
22time.sleep(10)
23query.stop()

You should see the data from Kafka properly parsed into a struct:

1+--------------------+
2|               value|
3+--------------------+
4|{18, Stranger Thi...|
5+--------------------+

Let's do some SQL on it now. We need to give a name to the temporary SQL view Spark is going to create from our DataFrame and use (we cannot use the name of the variable otherwise we'll run into a TABLE_OR_VIEW_NOT_FOUND error because it is not known from in the SQL context).

We execute this to create a temporary SQL view named netflix_view and display the average ratings by content:

1# we select value.* to avoid prefixing field names by "value", eg: value.name 
2json_df.select("value.*").createOrReplaceTempView("netflix_view")
3
4averageRatings = spark.sql("SELECT content, AVG(rating) FROM netflix_view GROUP BY content")
5query = averageRatings.writeStream.outputMode("complete").format("console").start()
6# we must switch to "complete" output mode because we GROUP BY otherwise Spark will fail with:
7# Append output mode not supported when there are streaming aggregations on streaming DataFrames/DataSets without watermar
8
9time.sleep(10)
10query.stop()

Each time we're going to produce data into our topic, we'll see the average ratings by content updated in real-time:

Let's say we produce this into our Kafka topic:

1{
2  "user_id": 112,
3  "content": "The Crown",
4  "watched_at": "2023-08-11 10:00:01",
5  "rating": 4
6}

Spark will display the update:

1+---------------+-----------+
2|        content|avg(rating)|
3+---------------+-----------+
4|Stranger Things|       3.75|
5|      The Crown|        4.0|
6+---------------+-----------+

Thanks to pyspark and Kafka, we have the power to run any SQL query on our data, and see the results updated in real-time in our terminal. Glorious!

Let's take a detour and explore alternatives using Conduktor to do SQL without Spark or without any 3rd party stream processing platform actually.

Faster: Use Console Kafka UI to do filtering

If we want to filter only fields (no GROUP BY), there is no need for pyspark, as Console is a powerful Kafka UI that provides various filtering mechanisms to filter data from Kafka.

Example below: we filter only the messages where rating > 4. We can add as many filters as we want, pretty cool and faster than writing a program.

Conduktor

Smarter: create VIEW over a topic (advanced)

If you're familiar with Conduktor Gateway, we can create a Virtual SQL Topic, a persistent SQL view over a topic. This is an alternative to Spark and Flink. There is no code to write, no stream processing framework to manage. It's totally seamless and costs absolutely $0 as everything is in memory and processed at runtime!

If you have deployed Conduktor Gateway, you can connect to it on Console and add the SQL Topic interceptor with your SQL query:

1{
2  "name": "my-virtual-sql-topic-interceptor",
3  "pluginClass": "io.conduktor.gateway.interceptor.VirtualSqlTopicPlugin",
4  "priority": 100,
5  "config": {
6    "virtualTopic": "good_ratings",
7    "statement": "SELECT * FROM my_netflix_topic WHERE rating >= 4"
8  }
9}

Using this, all users now have access to a topic named good_ratings (built from your SQL query). It's totally seamless that it's a virtual topic for all your users: it's not materialized (no storage, no partitions), but it acts like a normal topic. You can use it with kafka-console-consumer, Spring, etc.

A classic use-case we see it to hide the original topic my_netflix_topic to all users, and only expose a subset, like good_ratings here.

It's a great way to hide sensitive data, or to expose only the data you want to expose. You can also change the SQL query on-the-fly, without any consumer to be aware of it (if you want to refine your WHERE condition for instance).

Generating JSON with AI (advanced)# 

Ready to go mad and explore more options? Let's use OpenAI (ChatGPT) and Google to do some crazy stuff! This is an optional part as it is more advanced. Feel free to skip it if you're not interested.

We need an OpenAI account and a Google Cloud account. Register to both if not done yet.

When it's done, we need to:

Then export the following environment variables to make them available to pyspark (you have to quit pyspark (Ctrl+C), add these exports to your shell, and restart pyspark):

1# replace with your own values from OpenAI and Google. These values are fake.
2export OPENAI_API_KEY=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
3export GOOGLE_API_KEY=AIzaSyCfF5XZRSxkmUSjn5SE5tVBZLrzdyw2pFQ
4export GOOGLE_CSE_ID=5742144820bd0493e

We'll use pyspark-ai to work with OpenAI from pyspark.

  • If you have a paying OpenAPI subscription with access to ChatGPT-4, use this (default):
1from pyspark_ai import SparkAI
2spark_ai = SparkAI()
3spark_ai.activate()
  • If you don't have access to GPT4 ("The model gpt-4 does not exist or you do not have access to it"), use ChatGPT3.5 instead. For this, we use langchain to provide a specific model to SparkAI:
1from pyspark_ai import SparkAI
2from langchain.chat_models import ChatOpenAI
3llm = ChatOpenAI(model_name='gpt-3.5-turbo')
4spark_ai = SparkAI(llm=llm, verbose=True)
5spark_ai.activate()

The list of the available models are here: https://platform.openai.com/account/rate-limits, there are many models!

After the activation, we can use spark_ai to use the AI power. Let's start by creating a DataFrame from Wikipedia about the most subscribed Youtube channels:

1df = spark_ai.create_df("https://en.wikipedia.org/wiki/List_of_most-subscribed_YouTube_channels", [ "name", "subscribers", "language", "category", "country" ])
2df.show()

This will output the data properly extracted from Wikipedia:

1[..]
2+--------------------+-----------+--------+-------------+--------------------+
3|                name|subscribers|language|     category|             country|
4+--------------------+-----------+--------+-------------+--------------------+
5|            T-Series|      249.0|   Hindi|        Music|               India|
6|             MrBeast|      183.0| English|Entertainment|       United States|
7|           Cocomelon|      165.0| English|    Education|       United States|
8|Sony Entertainmen...|      162.0|   Hindi|Entertainment|               India|
9|     Kids Diana Show|      113.0| English|Entertainment|Ukraine- United S...|
10|           PewDiePie|      111.0| English|Entertainment|              Sweden|
11|         Like Nastya|      107.0| English|Entertainment|Russia- United St...|
12|       Vlad and Niki|      101.0| English|Entertainment|Russia- United St...|
13|   Zee Music Company|       99.5|   Hindi|        Music|               India|
14|                 WWE|       97.1| English|       Sports|       United States|
15|           Blackpink|       91.2|  Korean|        Music|         South Korea|
16|           Goldmines|       89.5|   Hindi|         Film|               India|
17|            Sony SAB|       85.2|   Hindi|Entertainment|               India|
18|     5-Minute Crafts|       80.2| English|       How-to|              Cyprus|
19|           BangtanTV|       76.4|  Korean|        Music|         South Korea|
20|         Hybe Labels|       72.6|  Korean|        Music|         South Korea|
21|              Zee TV|       72.4|   Hindi|Entertainment|               India|
22|       Justin Bieber|       71.9| English|        Music|              Canada|
23|            Pinkfong|       69.5| English|    Education|         South Korea|
24|ChuChu TV Nursery...|       67.5|   Hindi|    Education|               India|
25+--------------------+-----------+--------+-------------+--------------------+
26[..]

We are in 2023, right? Thanks to the advent of AI, we now have a user-friendly, natural language query system to manipulate data:

1>>> df.ai.verify("expect France not to be in the countries")
2Result: True
3# it hurts!
1>>>  df.ai.transform("per country").show()
2+--------------------+-----------------+
3|             country|total_subscribers|
4+--------------------+-----------------+
5|               India|           1312.2|
6|       United States|            678.5|
7|Ukraine- United S...|            113.0|
8|              Sweden|            111.0|
9|Russia- United St...|            208.0|
10|         South Korea|            309.7|
11|              Cyprus|             80.2|
12|              Canada|             71.9|
13|              Brazil|             66.6|
14|           Argentina|             59.7|
15+--------------------+-----------------+

Behind the scene, it's building a pyspark program (typically with SQL) to answer our queries, this is impressive.

If you're not using GPT-4, you might run into this error, because the dataset of the query is too big to be handled by GPT-3.5:

openai.error.InvalidRequestError: This model's maximum context length is 4097 tokens.
However, your messages resulted in 7337 tokens. Please reduce the length of the messages.

Another amazing thing is to ask a question to be plotted, and it will understand what we expect, build a pyspark program to make it happen, and open our browser for us:

1>>> df.ai.plot("by country")

Plot

Using AI to generate JSON into Kafka# 

Now that we understand how AI could help us, let's use it to generate random JSON based on our model.

We use directly ChatOpenAI to ask it to generate similar payloads from our model:

1from langchain.chat_models import ChatOpenAI
2from langchain.schema import *
3llm = ChatOpenAI(model_name='gpt-4') # or gpt-3.5 if you don't have a paying OpenAI subscription
4random_json = llm([HumanMessage(content="""generate random minified JSON payloads based on this: {
5  "user_id": 112,
6  "content": "The Crown",
7  "watched_at": "2023-08-11 10:00:01",
8  "rating": 4
9}""")]).content
10
11random_json

It takes a few seconds until ChatGPT understands and generates spot-on content:

1{"user_id":345,"content":"Breaking Bad","watched_at":"2022-12-14 18:30:12","rating":5}
2{"user_id":876,"content":"Friends","watched_at":"2023-01-03 21:00:00","rating":4}
3{"user_id":290,"content":"Stranger Things","watched_at":"2023-04-16 22:45:01","rating":5}
4{"user_id":789,"content":"The Witcher","watched_at":"2023-08-14 20:55:00","rating":4}
5{"user_id":654,"content":"Mandalorian","watched_at":"2023-07-20 19:10:10","rating":5}
6{"user_id":321,"content":"Peaky Blinders","watched_at":"2023-09-10 22:00:45","rating":4}
7{"user_id":154,"content":"Game of Thrones","watched_at":"2023-05-15 21:30:30","rating":3}
8{"user_id":903,"content":"Money Heist","watched_at":"2023-10-01 23:00:01","rating":4}
9{"user_id":567,"content":"Westworld","watched_at":"2023-06-12 20:10:10","rating":4}
10{"user_id":238,"content":"Better Call Saul","watched_at":"2023-07-11 19:00:00","rating":5}

We will use this to produce data into Kafka:

1chatgpt_records = random_json.strip().split("\n")
2chatgpt_df = spark.createDataFrame([Row(value=x) for x in chatgpt_records])
3chatgpt_df.selectExpr("CAST(value AS STRING)").write.format("kafka").options(**kafka_options).option("topic", "hello").save()

Each time we execute a new ChatGPT generation that we send into Kafka, we should see updated results in our pyspark terminal as our query averageRatings is still running (otherwise, rerun it):

1# Remember: averageRatings = spark.sql("SELECT content, AVG(rating) FROM netflix_view GROUP BY content")
2
3+---------------+-----------+
4|        content|avg(rating)|
5+---------------+-----------+
6|      Westworld|        4.0|
7|    Money Heist|        4.0|
8|Stranger Things|        4.0|
9|      The Crown|        4.0|
10|         Narcos|        5.0|
11|     The Office|        4.0|
12| Peaky Blinders|        5.0|
13|   Breaking Bad|        5.0|
14|The Mandalorian|        5.0|
15|        Friends|        3.0|
16|Game of Thrones|        4.0|
17+---------------+-----------+

The loop is closed! We produce random JSON data using AI into Kafka, and read it back in pyspark to do some SQL queries on it.

Conclusion# 

This is the end of this article as we have already covered a lot of things!

  • How to get started with PySpark and Kafka
  • How to do filtering and transformations using SQL
  • How to do the same simply with Conduktor
  • How to use OpenAI/ChatGPT to open up the potential of what we can do just by asking questions

pyspark is really useful when we need to run ad-hoc analysis locally, and explore our data using SQL. Combined to Conduktor with its powerful Kafka UI and Gateway to do SQL on Kafka, it's a powerful stack to explore real-time data and do some data exploration and data science for cheap.

We aim to accelerate Kafka projects delivery by making developers and organizations more efficient with Kafka.