How to Filter Rows Based on a Subquery Result in a PySpark DataFrame: The Ultimate Guide

Diving Straight into Filtering Rows with Subquery Results in a PySpark DataFrame

Filtering rows in a PySpark DataFrame is a cornerstone of data processing for data engineers working with Apache Spark in ETL pipelines, data cleaning, or analytics. Sometimes, you need to filter rows based on conditions derived from another query, known as a subquery. For example, you might want to select employees from departments with average salaries above a certain threshold. Subqueries allow you to dynamically define these conditions, making them essential for complex data workflows. This guide is tailored for data engineers with intermediate PySpark knowledge, building on your interest in PySpark filtering techniques [Timestamp: March 16, 2025]. If you’re new to PySpark, start with our PySpark Fundamentals.

We’ll cover the basics of subquery-based filtering, advanced subquery scenarios, handling nested data, SQL-based approaches, and optimizing performance. Each section includes practical code examples, outputs, and common pitfalls, explained in a clear, conversational tone to keep things actionable and relevant.

Understanding Subquery-Based Filtering in PySpark

A subquery is a query nested within another query, often used in a WHERE clause to define a dynamic condition. In PySpark, subqueries are typically executed using SQL expressions via spark.sql(), as the DataFrame API doesn’t directly support subquery syntax. By registering DataFrames as temporary views, you can use SQL subqueries to filter rows based on results from another dataset or aggregation, such as selecting rows matching IDs from a computed list.

Basic Subquery Filtering Example

Let’s filter employees from departments with at least one employee earning above $50,000.

Copy codefrom pyspark.sql import SparkSession

# Initialize Spark session
spark = SparkSession.builder.appName("SubqueryFilterExample").getOrCreate()

# Create employees DataFrame
employees_data = [
    (1, "Alice", 30, 50000, 101),
    (2, "Bob", 25, 45000, 102),
    (3, "Charlie", 35, 60000, 103),
    (4, "David", 28, 40000, 101),
    (5, "Eve", 32, 55000, 103)
]
employees = spark.createDataFrame(employees_data, ["employee_id", "name", "age", "salary", "dept_id"])

# Register DataFrame as a temporary view
employees.createOrReplaceTempView("employees")

# SQL query with subquery
filtered_df = spark.sql("""
    SELECT *
    FROM employees
    WHERE dept_id IN (
        SELECT dept_id
        FROM employees
        WHERE salary > 50000
    )
""")

# Show results
filtered_df.show()

# Output:
# +-----------+-------+---+------+-------+
# |employee_id|   name|age|salary|dept_id|
# +-----------+-------+---+------+-------+
# |          3|Charlie| 35| 60000|    103|
# |          5|    Eve| 32| 55000|    103|
# +-----------+-------+---+------+-------+

# Validate row count
assert filtered_df.count() == 2, "Expected 2 rows from department 103"

What’s Happening Here? The subquery SELECT dept_id FROM employees WHERE salary > 50000 identifies departments with at least one employee earning over $50,000 (dept_id 103). The outer query filters employees to include only rows where dept_id is in this list. This is a clean way to apply dynamic conditions based on aggregated or filtered data.

Key Methods:

• createOrReplaceTempView(viewName): Registers a DataFrame as a temporary view for SQL queries.
• spark.sql(sqlQuery): Executes an SQL query, including subqueries, and returns a DataFrame.
• IN clause: Filters rows where a column matches values from a subquery.

Common Mistake: Unregistered view.

Copy code# Incorrect: Querying unregistered view
spark.sql("SELECT * FROM employees WHERE dept_id IN (SELECT dept_id FROM employees WHERE salary > 50000)")  # Raises AnalysisException

# Fix: Register the view
employees.createOrReplaceTempView("employees")
spark.sql("SELECT * FROM employees WHERE dept_id IN (SELECT dept_id FROM employees WHERE salary > 50000)")

Error Output: AnalysisException: Table or view not found: employees.

Fix: Register the DataFrame as a view using createOrReplaceTempView() before running the query.

Advanced Subquery Filtering with Aggregations

Subqueries are particularly powerful when combined with aggregations, such as filtering rows based on group-level statistics. This allows you to apply conditions like selecting rows from groups meeting a threshold (e.g., departments with high average salaries).

Example: Filtering by Department Average Salary

Let’s filter employees from departments where the average salary exceeds $50,000.

Copy code# SQL query with aggregated subquery
filtered_df = spark.sql("""
    SELECT *
    FROM employees
    WHERE dept_id IN (
        SELECT dept_id
        FROM employees
        GROUP BY dept_id
        HAVING AVG(salary) > 50000
    )
""")

# Show results
filtered_df.show()

# Output:
# +-----------+-------+---+------+-------+
# |employee_id|   name|age|salary|dept_id|
# +-----------+-------+---+------+-------+
# |          3|Charlie| 35| 60000|    103|
# |          5|    Eve| 32| 55000|    103|
# +-----------+-------+---+------+-------+

# Validate
assert filtered_df.count() == 2, "Expected 2 rows from department 103"

What’s Going On? The subquery groups employees by dept_id, calculates the average salary per department, and keeps departments where the average exceeds $50,000 (dept_id 103). The outer query filters employees to include only rows from these departments. This is perfect for scenarios where you need to filter based on group-level metrics, like identifying high-paying departments.

Common Mistake: Missing HAVING clause for aggregations.

Copy code# Incorrect: Using WHERE for group condition
spark.sql("""
    SELECT *
    FROM employees
    WHERE dept_id IN (
        SELECT dept_id
        FROM employees
        GROUP BY dept_id
        WHERE AVG(salary) > 50000
    )
""")  # Raises SyntaxError

# Fix: Use HAVING
spark.sql("""
    SELECT *
    FROM employees
    WHERE dept_id IN (
        SELECT dept_id
        FROM employees
        GROUP BY dept_id
        HAVING AVG(salary) > 50000
    )
""")

Error Output: SyntaxError: cannot use WHERE after GROUP BY; use HAVING.

Fix: Use HAVING to filter groups after aggregation in subqueries.

Filtering Nested Data with Subqueries

Nested data, like structs, is common in semi-structured datasets. Subqueries can filter rows based on nested fields by accessing them with dot notation, making them versatile for complex data.

Example: Filtering by Nested Email Subquery

Suppose employees has a contact struct with email and phone. We want to filter employees from departments where at least one employee has a corporate email.

Copy codefrom pyspark.sql.types import StructType, StructField, StringType, IntegerType

# Define schema with nested struct
schema = StructType([
    StructField("employee_id", IntegerType()),
    StructField("name", StringType()),
    StructField("contact", StructType([
        StructField("email", StringType()),
        StructField("phone", StringType())
    ])),
    StructField("dept_id", IntegerType())
])

# Create employees DataFrame
employees_data = [
    (1, "Alice", {"email": "alice@company.com", "phone": "123-456-7890"}, 101),
    (2, "Bob", {"email": "bob@other.com", "phone": "234-567-8901"}, 102),
    (3, "Charlie", {"email": "charlie@company.com", "phone": "345-678-9012"}, 103),
    (4, "David", {"email": "david@gmail.com", "phone": "456-789-0123"}, 101)
]
employees = spark.createDataFrame(employees_data, schema)

# Register DataFrame as a temporary view
employees.createOrReplaceTempView("employees")

# SQL query with subquery on nested field
filtered_df = spark.sql("""
    SELECT *
    FROM employees
    WHERE dept_id IN (
        SELECT dept_id
        FROM employees
        WHERE contact.email LIKE '%company.com'
    )
""")

# Show results
filtered_df.show()

# Output:
# +-----------+-----+--------------------+-------+
# |employee_id| name|             contact|dept_id|
# +-----------+-----+--------------------+-------+
# |          1|Alice|{alice@company.co...|    101|
# |          3|Charlie|{charlie@company....|    103|
# |          4| David|{david@gmail.com,...|    101|
# +-----------+-----+--------------------+-------+

# Validate
assert filtered_df.count() == 3, "Expected 3 rows from departments 101 and 103"

What’s Going On? The subquery selects dept_id values where contact.email ends with "@company.com" (dept_ids 101 and 103). The outer query filters employees to include all rows from these departments, even if some employees (e.g., David) don’t have a corporate email. This is useful for JSON-like data where nested fields drive filtering logic, aligning with your interest in complex data structures [Timestamp: March 27, 2025].

Common Mistake: Incorrect nested field reference.

Copy code# Incorrect: Non-existent field
spark.sql("""
    SELECT *
    FROM employees
    WHERE dept_id IN (
        SELECT dept_id
        FROM employees
        WHERE contact.mail LIKE '%company.com'
    )
""")  # Raises AnalysisException

# Fix: Verify schema
employees.printSchema()
spark.sql("""
    SELECT *
    FROM employees
    WHERE dept_id IN (
        SELECT dept_id
        FROM employees
        WHERE contact.email LIKE '%company.com'
    )
""")

Error Output: AnalysisException: cannot resolve 'contact.mail'.

Fix: Use printSchema() to confirm nested field names.

Subquery Filtering with SQL Expressions

Since subqueries are inherently SQL-based in PySpark, the SQL approach is the primary method. However, you can combine subqueries with DataFrame operations for hybrid workflows, using spark.sql() for the subquery and DataFrame methods for further processing.

Example: Hybrid Subquery and DataFrame Filtering

Let’s filter employees from high-salary departments and then apply an additional DataFrame filter for age.

Copy code# SQL subquery for high-salary departments
high_salary_depts = spark.sql("""
    SELECT dept_id
    FROM employees
    GROUP BY dept_id
    HAVING AVG(salary) > 50000
""").collect()
dept_ids = [row["dept_id"] for row in high_salary_depts]

# DataFrame filter with subquery result
filtered_df = employees.filter(col("dept_id").isin(dept_ids)) \
                       .filter(col("age") > 30)

# Show results
filtered_df.show()

# Output:
# +-----------+-------+---+------+-------+
# |employee_id|   name|age|salary|dept_id|
# +-----------+-------+---+------+-------+
# |          3|Charlie| 35| 60000|    103|
# +-----------+-------+---+------+-------+

# Validate
assert filtered_df.count() == 1

What’s Going On? The subquery identifies departments with average salaries above $50,000 (dept_id 103). We collect the dept_id values into a Python list and use isin() in a DataFrame filter, combined with an age condition (age > 30). This hybrid approach leverages SQL for complex subqueries and DataFrame methods for additional flexibility.

Common Mistake: Large subquery result in isin().

Copy code# Inefficient: Large list in isin
large_dept_list = [i for i in range(1000000)]  # Large list causes performance issues
filtered_df = employees.filter(col("dept_id").isin(large_dept_list))

# Fix: Use join for large subquery results
high_salary_depts = spark.sql("SELECT dept_id FROM employees GROUP BY dept_id HAVING AVG(salary) > 50000")
filtered_df = employees.join(high_salary_depts, "dept_id", "inner")

Error Output: No error, but slow performance due to broadcasting a large list.

Fix: Use a join for large subquery results to leverage Spark’s distributed processing.

Optimizing Subquery-Based Filtering Performance

Subqueries can be resource-intensive, especially with large datasets or complex aggregations. Here are four strategies to optimize performance, drawing on your interest in Spark optimization [Timestamp: March 19, 2025].

1. Select Relevant Columns: Include only necessary columns in both the subquery and outer query to reduce shuffling.
2. Filter Early: Apply preliminary filters before subqueries to minimize the dataset.
3. Partition Data: Partition by frequently filtered columns (e.g., dept_id) for faster queries.
4. Cache Results: Cache subquery results or filtered DataFrames for reuse.

Example: Optimized Subquery Filtering

Copy code# Filter early and select relevant columns
optimized_df = employees.select("employee_id", "name", "dept_id", "salary") \
                       .filter(col("dept_id").isNotNull())

# Register filtered DataFrame
optimized_df.createOrReplaceTempView("filtered_employees")

# Optimized SQL query with subquery
filtered_df = spark.sql("""
    SELECT employee_id, name, dept_id, salary
    FROM filtered_employees
    WHERE dept_id IN (
        SELECT dept_id
        FROM filtered_employees
        GROUP BY dept_id
        HAVING AVG(salary) > 50000
    )
""").cache()

# Show results
filtered_df.show()

# Output:
# +-----------+-------+-------+------+
# |employee_id|   name|dept_id|salary|
# +-----------+-------+-------+------+
# |          3|Charlie|    103| 60000|
# |          5|    Eve|    103| 55000|
# +-----------+-------+-------+------+

# Validate
assert filtered_df.count() == 2

What’s Happening? We filter out null dept_id values and select only employee_id, name, dept_id, and salary to reduce data processed. The subquery identifies high-salary departments, and the outer query filters accordingly. Caching the result ensures efficiency for downstream operations, aligning with your focus on efficient ETL pipelines [Timestamp: March 15, 2025].

Wrapping Up Your Subquery-Based Filtering Mastery

Filtering PySpark DataFrame rows based on subquery results is a powerful technique for dynamic, data-driven conditions. From basic IN subqueries to aggregated conditions, nested data, hybrid SQL-DataFrame workflows, and performance optimizations, you’ve got a robust toolkit for complex filtering tasks. Try these methods in your next Spark project and share your insights on X. For more DataFrame operations, check out DataFrame Transformations.

More Spark Resources to Keep You Going

• Apache Spark Documentation
• Databricks Spark Guide
• PySpark DataFrame Basics
• PySpark Performance Tuning

Published: April 17, 2025