Ubuntu – polars: efficient way to apply function to filter column of strings

Let’s see if I can help with the performance concerns. There may be a way to get your algorithm to run in a performance envelope that is acceptable for your system.

For the benchmarking below, I created the following dataset of over a billion records, along with some extra columns (to simulate data other than the string column that is being processed).

shape: (1_073_741_824, 6)
┌───────────────────────────────────┬────────┬───────┬─────┬─────────────────────┬─────────────────────┐
│ foo                               ┆ string ┆ float ┆ int ┆ datetime            ┆ date                │
│ ---                               ┆ ---    ┆ ---   ┆ --- ┆ ---                 ┆ ---                 │
│ str                               ┆ str    ┆ f32   ┆ i32 ┆ datetime[μs]        ┆ datetime[μs]        │
╞═══════════════════════════════════╪════════╪═══════╪═════╪═════════════════════╪═════════════════════╡
│ COOOPS.autom.SAPF124              ┆ string ┆ 1.0   ┆ 1   ┆ 2024-03-01 00:00:00 ┆ 2024-03-01 00:00:00 │
│ OSS REEE PAAA comp. BEEE  atm 60… ┆ string ┆ 1.0   ┆ 1   ┆ 2024-03-01 00:00:00 ┆ 2024-03-01 00:00:00 │
│ ABCD 600000000397/7667896-6/REG.… ┆ string ┆ 1.0   ┆ 1   ┆ 2024-03-01 00:00:00 ┆ 2024-03-01 00:00:00 │
│ OSS REFF pagopago cost. Becf  at… ┆ string ┆ 1.0   ┆ 1   ┆ 2024-03-01 00:00:00 ┆ 2024-03-01 00:00:00 │
│ COOOPS.autom.SAPF124              ┆ string ┆ 1.0   ┆ 1   ┆ 2024-03-01 00:00:00 ┆ 2024-03-01 00:00:00 │
│ …                                 ┆ …      ┆ …     ┆ …   ┆ …                   ┆ …                   │
│ OSS REFF pagopago cost. Becf  at… ┆ string ┆ 1.0   ┆ 1   ┆ 2024-03-01 00:00:00 ┆ 2024-03-01 00:00:00 │
│ COOOPS.autom.SAPF124              ┆ string ┆ 1.0   ┆ 1   ┆ 2024-03-01 00:00:00 ┆ 2024-03-01 00:00:00 │
│ OSS REEE PAAA comp. BEEE  atm 60… ┆ string ┆ 1.0   ┆ 1   ┆ 2024-03-01 00:00:00 ┆ 2024-03-01 00:00:00 │
│ ABCD 600000000397/7667896-6/REG.… ┆ string ┆ 1.0   ┆ 1   ┆ 2024-03-01 00:00:00 ┆ 2024-03-01 00:00:00 │
│ OSS REFF pagopago cost. Becf  at… ┆ string ┆ 1.0   ┆ 1   ┆ 2024-03-01 00:00:00 ┆ 2024-03-01 00:00:00 │
└───────────────────────────────────┴────────┴───────┴─────┴─────────────────────┴─────────────────────┘

To simulate the situation you describe, I saved the dataframe to a parquet file with default compression. This parquet file is then used as the input below (input_filepath in the code below).

For reference, the above dataset consumes 77 GB of RAM when fully loaded in RAM (per Polars estimated_size method). I’m running these benchmarks on a 32-core Threadripper Pro with 504 GB of available RAM.

Using collect

Using the normal collect method with default options, I ran the following code copied from @Hericks excellent answer above:

from time import perf_counter


start = perf_counter()
word_list = ["cost", "atm"]

num_dec_expr = pl.element().str.count_matches(r"[0-9_/]").cast(pl.Int32)
num_letter_expr = pl.element().str.count_matches(r"[A-Za-z]").cast(pl.Int32)
ratio_expr = (num_letter_expr - num_dec_expr) / (num_letter_expr + num_dec_expr)

(
    pl.scan_parquet(input_filepath)
    .with_columns(
        pl.col("foo")
        .str.to_lowercase()
        .str.replace_all(r"[^a-z0-9s_-%]", " ")
        .str.split(" ")
        .list.eval(
            pl.element().filter(
                pl.element().str.len_chars() > 0,
                pl.element().is_in(word_list)
                | num_letter_expr.eq(0)
                | ratio_expr.is_between(-0.7, 0),
            )
        )
        .list.join(" ")
        .alias("foo_clean")
    )
    .collect()
)

print("Elapsed time: ", perf_counter() - start)

The runtime performance was as you describe. The algorithm quickly reverts to single-threaded behavior, and allocates an inordinate amount of RAM (250 GB of RAM) per the top command. I killed the process after waiting for over 15 minutes.

Using collect(streaming=True, comm_subplan_elim=False)

By substituting the collect statement with collect(streaming=True, comm_subplan_elim=False), the above completes in a mere 167 seconds. And top shows that the algorithm pushes all 64 logical cores of my Threadripper Pro system to 100%.

However, the algorithm does continue to consume a large amount of RAM: 160 GB while running, just eyeballing the top command. Still, that certainly is better than the 250 GB of RAM when using collect with default options.

using sink_parquet

By substituting the sink_parquet method instead of the collect method and thus saving the results directly to disk, the algorithm ran in 459 seconds (about 7.5 minutes). And the RAM topped out at a mere 8GB while processing the file.

One thing I’d caution, the runtime behavior showed that it used all 64 logical cores of my system, but not to 100%. This could be due to an I/O bottleneck, but I doubt it. (My system has four Gen4 1TB NVME sticks in RAID 0 as a working project storage.) As such, I would suggest that using sink_parquet could take longer on a system where I/O bottlenecks are more pronounced.

Admittedly, streaming the results directly to disk may not be want you want. But it may be what you need to get past this step of your processing with acceptable run-time and memory footprints.

If any of this helps, well and good. And please continue to give credit for the accepted answer to @Hericks. His answer was spot-on regarding how to use Polars’ native string-manipulation capabilities.