benchmark filtering list in polars series

benchmark_filters.py

import polars as pl
import numpy as np
from itertools import product
import time
from datetime import datetime
import json


def gen_long_string(str_len=10, n_rows=10_000_000):
    rng = np.random.default_rng()
    return rng.integers(low=96, high=122, size=n_rows * str_len, dtype="uint32").view(
        f"U{str_len}"
    )


def run_benches(n, j):
    print(f"{datetime.now()}: n={n}, j={j}")
    remove_str = gen_long_string(6, 1)[0]
    df = (
        pl.concat(
            [
                pl.DataFrame(
                    {
                        "i": np.random.randint(0, n, n * j),
                        "a": gen_long_string(5, n * j),
                    }
                ),
                pl.select(i=pl.arange(n), a=pl.lit(remove_str)),
            ]
        )
        .sample(n * (j + 1), shuffle=True)
        .filter(pl.len().over("i") > 2)
    )

    s = df.group_by("i").agg("a").get_column("a")

    result = {}
    result["list_eval"] = {}
    strt = time.time()
    result["list_eval"]["answer"] = s.list.eval(
        pl.element().filter(pl.element() != remove_str)
    )
    result["list_eval"]["timing"] = time.time() - strt

    if n <= 1_000_000:
        result["map_elem"] = {}
        strt = time.time()
        result["map_elem"]["answer"] = s.map_elements(
            lambda x: [y for y in x if y != remove_str]
        )
        result["map_elem"]["timing"] = time.time() - strt

    result["gather"] = {}
    strt = time.time()
    result["gather"]["answer"] = (
        pl.select(a=s)
        .with_row_index("i")
        .group_by("i")
        .agg(
            pl.col("a")
            .list.gather(pl.arg_where(pl.col("a").explode() != remove_str))
            .first()
        )
        .get_column("a")
    )
    result["gather"]["timing"] = time.time() - strt

    result["filter"] = {}
    strt = time.time()
    result["filter"]["answer"] = (
        pl.select(a=s)
        .with_row_index("i")
        .explode("a")
        .filter(pl.col("a") != remove_str)
        .group_by("i")
        .agg("a")
        .get_column("a")
    )
    result["filter"]["timing"] = time.time() - strt

    result["set_diff"] = {}
    strt = time.time()
    result["set_diff"]["answer"] = s.list.set_difference([remove_str])
    result["set_diff"]["timing"] = time.time() - strt

    res_keys = list(result.keys())
    for i, key in enumerate(res_keys):
        for jkey in res_keys[i + 1 :]:
            if not result[key]["answer"].equals(result[jkey]["answer"]):
                result["inconsistent"] = (key, jkey)
    for key in res_keys:
        del result[key]["answer"]
    result["n"] = n
    result["j"] = j
    with open(f"./benchmarks/n{n}_j{j}.json", "w") as ff:
        ff.write(json.dumps(result))


def main():
    for nexp, j in product(range(1, 6), range(5, 55, 5)):
        n = int(10**nexp)
        run_benches(n, j)
    for n, j in product(range(200_000, 1_100_000, 100_000), range(5, 55, 5)):
        run_benches(n, j)
    for n, j in product(range(1_500_000, 5_500_000, 500_000), range(5, 55, 5)):
        run_benches(n, j)
        
## I just run this line by line in VSC  
def interactive():
    import os
    import plotly.express as px
    results = pl.concat(
        [pl.read_json(f"./benchmarks/{x}") for x in os.listdir("./benchmarks")],
        how='diagonal_relaxed'
    )
    results = results.select(
        pl.col(x).struct.field("timing").alias(x)
        if isinstance(y, pl.Struct)
        else pl.col(x)
        for x, y in results.schema.items()
    ).drop("inconsistent").sort(['n','j'])
    px.line(
        results.filter(j=50),
        x='n',
        y=['list_eval','filter','set_diff'],
    )
    
    px.line(
    results.filter(j=50),
    x='n',
    y=['map_elem','gather','list_eval','filter','set_diff'],
    log_y=True
)



if __name__ == "__main__":
    main()

        

results.csv