LazyFrames in Polars come with powerful optimization capabilities, and the explain() method gives us visibility into how these optimizations work. By examining both optimized and unoptimized query plans, we can better understand how Polars improves our code’s performance and write more efficient data processing pipelines.
Basic Query Optimization
Let’s start with a simple example that demonstrates how Polars optimizes basic filtering operations:
import polars as pl
# Create sample data
df = pl.DataFrame({
'product': ['A', 'B', 'C', 'A', 'B', 'C'],
'quantity': [10, 20, 15, 25, 30, 35],
'price': [100, 200, 150, 250, 300, 350]
})
# Create a lazy query with multiple filters
query = (df.lazy()
.filter(pl.col('quantity') > 15)
.filter(pl.col('price') > 200)
.select(['product', 'quantity', 'price']))
# Show unoptimized plan
print("Unoptimized Plan:")
print(query.explain(optimized=False))
# Show optimized plan
print("\nOptimized Plan:")
print(query.explain()) # optimized=True is default
When we run this code, we get:
Unoptimized Plan:
SELECT [col("product"), col("quantity"), col("price")] FROM
FILTER [(col("price")) > (200)] FROM
FILTER [(col("quantity")) > (15)] FROM
DF ["product", "quantity", "price"]; PROJECT */3 COLUMNS; SELECTION: None
Optimized Plan:
DF ["product", "quantity", "price"]; PROJECT 3/3 COLUMNS; SELECTION: [([(col("quantity")) > (15)]) & ([(col("price")) > (200)])]
Notice how the unoptimized plan shows three separate operations read from bottom to top: first filtering by quantity (>15), then filtering by price (>200), and finally selecting the columns. In contrast, the optimized plan combines both filter conditions into a single selection operation, reducing the number of passes through the data. This optimization becomes particularly important when working with large datasets.
Working with Aggregations
Let’s examine a more complex example involving grouping and aggregations:
# Create a query with grouping and multiple aggregations
complex_query = (df.lazy()
.group_by('product')
.agg([
pl.col('quantity').sum().alias('total_quantity'),
pl.col('price').mean().alias('avg_price')
])
.filter(pl.col('total_quantity') > 30)
.sort('avg_price', descending=True))
# Compare plans
print("Unoptimized Plan:")
print(complex_query.explain(optimized=False))
print("\nOptimized Plan:")
print(complex_query.explain())
Running this code produces:
Unoptimized Plan:
SORT BY [col("avg_price")]
FILTER [(col("total_quantity")) > (30)] FROM
AGGREGATE
[col("quantity").sum().alias("total_quantity"), col("price").mean().alias("avg_price")] BY [col("product")] FROM
DF ["product", "quantity", "price"]; PROJECT */3 COLUMNS; SELECTION: None
Optimized Plan:
SORT BY [col("avg_price")]
FILTER [(col("total_quantity")) > (30)] FROM
AGGREGATE
[col("quantity").sum().alias("total_quantity"), col("price").mean().alias("avg_price")] BY [col("product")] FROM
DF ["product", "quantity", "price"]; PROJECT 3/3 COLUMNS; SELECTION: None
In this more complex example, we can see that both plans maintain the same basic structure. However, the optimized plan specifies that it’s using all columns (3/3) rather than using a wildcard (*). This clarity in column projection can help Polars make better decisions about memory allocation and data movement.
Let’s look at another aggregation example that demonstrates more substantial optimizations:
# Let's create a more complex example that shows clearer optimization
complex_query_2 = (df.lazy()
.filter(pl.col('price') > 200)
.group_by('product')
.agg([
pl.col('quantity').sum().alias('total_quantity'),
pl.col('price').mean().alias('avg_price')
])
.filter(pl.col('total_quantity') > 30)
.filter(pl.col('avg_price') > 275)
.sort('avg_price', descending=True))
print("Unoptimized Plan:")
print(complex_query_2.explain(optimized=False))
print("\nOptimized Plan:")
print(complex_query_2.explain())
This query produces a more interesting optimization pattern:
Unoptimized Plan:
SORT BY [col("avg_price")]
FILTER [(col("avg_price")) > (275.0)] FROM
FILTER [(col("total_quantity")) > (30)] FROM
AGGREGATE
[col("quantity").sum().alias("total_quantity"), col("price").mean().alias("avg_price")] BY [col("product")] FROM
FILTER [(col("price")) > (200)] FROM
DF ["product", "quantity", "price"]; PROJECT */3 COLUMNS; SELECTION: None
Optimized Plan:
SORT BY [col("avg_price")]
FILTER [([(col("total_quantity")) > (30)]) & ([(col("avg_price")) > (275.0)])] FROM
AGGREGATE
[col("quantity").sum().alias("total_quantity"), col("price").mean().alias("avg_price")] BY [col("product")] FROM
DF ["product", "quantity", "price"]; PROJECT 3/3 COLUMNS; SELECTION: [(col("price")) > (200)]
In this more complex example, we can observe two significant optimizations. First, Polars combines the two post-aggregation filters (total_quantity > 30 and avg_price > 275) into a single filter operation using a logical AND. Second, it moves the initial price filter (price > 200) to the selection phase, demonstrating how Polars can reorganize operations for better efficiency. These optimizations reduce the number of intermediate steps and potentially improve memory usage by minimizing the data processed at each stage.
Key Optimization Principles
The explain() method helps us understand query optimization in several ways:
- It shows how Polars combines multiple filter conditions into a single operation
- It provides clarity about which columns are being projected and used
- It reveals the sequence of operations that will be performed on our data
By default, explain() shows the optimized plan (optimized=True), but comparing it with the unoptimized version (optimized=False) can help us understand the improvements Polars makes to our queries. This knowledge is particularly valuable when working with large datasets where optimization can significantly impact performance.
Polars’ query optimizer demonstrates sophisticated capabilities in restructuring our operations for better performance. Understanding these optimization patterns helps us write more thoughtful queries that align with Polars’ strengths, particularly when working with complex transformations on large datasets. By leveraging insights from the explain() method, we can design data processing pipelines that take full advantage of Polars’ optimization capabilities from the outset.
