如何基于元组列表在PySpark DataFrame中筛选行？

Great question! Filtering rows based on a list of tuples (like matching combinations of aid and bid in your example) is a common task in PySpark, and there are a few efficient ways to do it depending on your use case. Let's walk through the most practical methods with your sample DataFrame.

First, let's set up our sample data and target tuple list to make the examples concrete:

from pyspark.sql import SparkSession

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

# Your sample DataFrame data
data = [
    (1, 1, 81.0, "2006-08-25 14:13:00"),
    (1, 1, 81.0, "2006-08-25 14:27:00"),
    (1, 2, 81.0, "2006-08-25 14:56:00"),
    (1, 2, 81.0, "2006-08-25 15:00:00"),
    (1, 3, 81.0, "2006-08-25 15:31:00"),
    (1, 3, 81.0, "2006-08-25 15:38:00"),
    (1, 4, 0.0, "2006-08-30 11:59:00"),
    (1, 4, 0.0, "2006-08-30 13:59:00"),
    (2, 1, 0.0, "2006-08-30 12:11:00"),
    (2, 1, 0.0, "2006-08-30 13:11:00")
]

df = spark.createDataFrame(data, ["aid", "bid", "value", "time"])

# Let's say we want to keep rows where (aid, bid) is in this list
target_tuples = [(1, 2), (2, 1)]

Method 1: Use isin() with Column Tuples (Simplest for Small/Medium Lists) #

If your target tuple list isn't extremely large (hundreds or low thousands of elements), this is the most straightforward approach. PySpark 2.3+ supports checking if a tuple of columns is in a list of tuples directly:

filtered_df = df.filter((df.aid, df.bid).isin(target_tuples))
filtered_df.show()

Output: #

+---+---+-----+--------------------+
|aid|bid|value|                time|
+---+---+-----+--------------------+
|  1|  2| 81.0|2006-08-25 14:56:...|
|  1|  2| 81.0|2006-08-25 15:00:...|
|  2|  1|  0.0|2006-08-30 12:11:...|
|  2|  1|  0.0|2006-08-30 13:11:...|
+---+---+-----+--------------------+

Pros: Quick to write, no extra steps.
Cons: Not ideal for very large tuple lists (can cause performance issues due to long filter clauses).

Method 2: Join with a Temporary DataFrame (Best for Large Lists) #

If your target tuple list is large (thousands+ elements), using a join is more efficient. Spark can optimize this operation, especially if you broadcast the small target DataFrame to avoid shuffling large data:

from pyspark.sql.functions import broadcast

# Create a small DataFrame from our target tuples
target_df = spark.createDataFrame(target_tuples, ["aid", "bid"])

# Use inner join to keep only matching rows (broadcast hint speeds up the join)
filtered_df = df.join(broadcast(target_df), on=["aid", "bid"], how="inner")
filtered_df.show()

Pros: Scalable for large datasets and large tuple lists; Spark optimizes the join efficiently.
Cons: Slightly more code than isin(), but worth it for performance at scale.

Method 3: Use SQL Syntax (If You Prefer SQL) #

If you're more comfortable writing SQL queries, you can register your DataFrame as a temporary view and use an IN clause with tuples:

# Register the original DataFrame as a temporary SQL view
df.createOrReplaceTempView("my_table")

# Convert the target tuples to a string format suitable for SQL
tuple_str = ", ".join([f"({a}, {b})" for a, b in target_tuples])

# Run the SQL query
filtered_df = spark.sql(f"SELECT * FROM my_table WHERE (aid, bid) IN ({tuple_str})")
filtered_df.show()

Pros: Familiar if you know SQL; useful for complex queries combined with filtering.
Cons: Requires string formatting, which can be error-prone if tuples contain special characters (though not an issue here with integers).

Key Notes: #

• Always choose the method based on the size of your tuple list: isin() for small lists, joins for large ones.
• The broadcast hint in Method 2 is optional but highly recommended when the target DataFrame is small—it sends the target data to all worker nodes, avoiding expensive data shuffles.

内容的提问来源于stack exchange，提问作者Stefan Falk