Optimization and query plans

Scaling and Optimizing Data Pipelines with Polars

Liam Brannigan

Data Scientist & Polars Contributor

Introducing query optimization

Image of table with rows and columns

Introducing query optimization

Image of table with rows and columns

Introducing query optimization

Image of running operations in serial

Introducing query optimization

Image of running operations in parallel

Introducing query optimization

Pipeline with duplicate operations carried out twice

Introducing query optimization

Pipeline with duplicate operations carried out twice

Most request types by department

department_request_types = (
    requests





)

Most request types by department

department_request_types = (
    requests
    .filter(pl.col("STATUS") == "Completed")




)

Most request types by department

department_request_types = (
    requests
    .filter(pl.col("STATUS") == "Completed")
    .group_by("DEPARTMENT")



)

Most request types by department

department_request_types = (
    requests    
    .filter(pl.col("STATUS") == "Completed")
    .group_by("DEPARTMENT")
    .agg(pl.col("TYPE").n_unique().alias("n_request_types"))


)

Most request types by department

department_request_types = (
    requests
    .filter(pl.col("STATUS") == "Completed")
    .group_by("DEPARTMENT")
    .agg(pl.col("TYPE").n_unique().alias("n_request_types"))
    .sort("n_request_types", descending=True)
    .head(5)
)

Naive plan

Unoptimized plan

print(department_request_types)

Unoptimized plan

print(department_request_types)








        Csv SCAN [311_Service_Requests.csv]
        PROJECT */39 COLUMNS

Unoptimized plan

print(department_request_types)






      FILTER [(col("STATUS")) == ("Completed")]
      FROM
        Csv SCAN [311_Service_Requests.csv]
        PROJECT */39 COLUMNS

Unoptimized plan

print(department_request_types)



    AGGREGATE[maintain_order: false]
      [col("TYPE").n_unique().alias("n_request_types")] BY [col("DEPARTMENT")]
      FROM
      FILTER [(col("STATUS")) == ("Completed")]
      FROM
        Csv SCAN [311_Service_Requests.csv]
        PROJECT */39 COLUMNS

Unoptimized plan

print(department_request_types)

SLICE[offset: 0, len: 5]
  SORT BY [descending: [true]] [col("n_request_types")]
    AGGREGATE[maintain_order: false]
      [col("TYPE").n_unique().alias("n_request_types")] BY [col("DEPARTMENT")]
      FROM
      FILTER [(col("STATUS")) == ("Completed")]
      FROM
        Csv SCAN [311_Service_Requests.csv]
        PROJECT */39 COLUMNS

Optimized plan

print(department_request_types.explain())

Optimized plan







        Csv SCAN [311_Service_Requests.csv]
        PROJECT 3/39 COLUMNS
        SELECTION: [(col("STATUS")) == ("Completed")]

Optimized plan






      simple pi 2/2 ["TYPE", "DEPARTMENT"]
        Csv SCAN [311_Service_Requests.csv]
        PROJECT 3/39 COLUMNS
        SELECTION: [(col("STATUS")) == ("Completed")]

Optimized plan



    AGGREGATE[maintain_order: false]
      [col("TYPE").n_unique().alias("n_request_types")] BY [col("DEPARTMENT")]
      FROM
      simple pi 2/2 ["TYPE", "DEPARTMENT"]
        Csv SCAN [311_Service_Requests.csv]
        PROJECT 3/39 COLUMNS
        SELECTION: [(col("STATUS")) == ("Completed")]

Optimized plan

SORT BY [slice: (0, 10, ...), descending: [true]] [col("n_request_types")]
  FILTER col("n_request_types").dynamic_predicate() FROM
    AGGREGATE[maintain_order: false]
      [col("TYPE").n_unique().alias("n_request_types")] BY [col("DEPARTMENT")]
      FROM
      simple pi 2/2 ["TYPE", "DEPARTMENT"]
        Csv SCAN [311_Service_Requests.csv]
        PROJECT 3/39 COLUMNS
        SELECTION: [(col("STATUS")) == ("Completed")]

Find the top 10 rows in n_request_types
Filter out the rest
Sort the 10 rows

Optimized plan as a graph

print(department_request_types.show_graph())

Graph view of the optimized plan showing CSV scan with optimizations.

Further optimizations

(
    requests








)

Further optimizations

(
    requests
    .filter(pl.col("STATUS") == "Completed")
    .filter(pl.col("DEPARTMENT") == "Sanitation")






)

Further optimizations

(
    requests
    .filter(pl.col("STATUS") == "Completed")
    .filter(pl.col("DEPARTMENT") == "Sanitation")
    .with_columns(
        pl.col("TYPE").str.to_lowercase().alias("type_lower")
    )
    .with_columns(
        pl.col("STATUS").str.to_lowercase().alias("status_lower")
    )
)

Further optimizations



  Csv SCAN [311_Service_Requests.csv]
  PROJECT */39 COLUMNS
  SELECTION: [([(col("DEPARTMENT")) == ("Sanitation")]) & ([(col("STATUS")) == ("Completed")])]

Combined AND predicate

Further optimizations

 WITH_COLUMNS:
 [col("TYPE").str.to_lowercase().alias("type_lower"), col("STATUS").str.to_lowercase().alias("status_lower")]
  Csv SCAN [311_Service_Requests.csv]
  PROJECT */39 COLUMNS
  SELECTION: [([(col("DEPARTMENT")) == ("Sanitation")]) & ([(col("STATUS")) == ("Completed")])]

Combined AND predicate
Clustered WITH_COLUMNS expressions

Let's practice!

Scaling and Optimizing Data Pipelines with Polars