Performance issues with group by

[billylanchantin]

@mhanberg Noted alerted us to some unexpectedly slow performance in Discord. I'm making this as a placeholder for myself so I don't forget about the investigation.

Benchmarking

Elixir code & results

defmodule MotchBenchmarkTest do
  use ExUnit.Case, async: true
  use ExUnitProperties

  require Explorer.DataFrame, as: DataFrame

  setup do
    n_rows = 750_000
    results = ["overbaked", "underbaked", "justright"]

    data_frame =
      DataFrame.new(
        alice_id: Enum.map(1..n_rows, fn _ -> :rand.uniform(5) end),
        bob_id: Enum.map(1..n_rows, fn _ -> :rand.uniform(10) end),
        carol_id: Enum.map(1..n_rows, fn _ -> :rand.uniform(15) end),
        dave_id: Enum.map(1..n_rows, fn _ -> :rand.uniform(20) end),
        foo_id: Enum.map(1..n_rows, fn _ -> :rand.uniform(25) end),
        ingredients: Enum.map(1..n_rows, fn _ -> :rand.uniform(30) end),
        sim_id: Enum.map(1..n_rows, fn _ -> :rand.uniform(35) end),
        result: Enum.flat_map(1..n_rows, fn _ -> Enum.take_random(results, 1) end)
      )

    DataFrame.to_parquet(data_frame, "motch.parquet")

    %{data_frame: data_frame}
  end

  test "motch", %{data_frame: data_frame} do
    cookies = do_cookies(data_frame)
    rest = do_rest(data_frame)

    # Benchee.run(
    #   %{
    #     "1_cookies" => fn ->
    #       do_cookies(data_frame)
    #     end,
    #     "2_rest" => fn ->
    #       do_rest(data_frame)
    #     end,
    #     "3_join" => fn ->
    #       DataFrame.join(rest, cookies, on: [:sim_id, :foo_id], how: :left)
    #     end
    #   },
    #   time: 10,
    #   memory_time: 2
    # )
    # |> IO.inspect()
  end

  def do_cookies(data_frame) do
    cols = [
      :alice_id,
      :bob_id,
      :carol_id,
      :dave_id
    ]

    for col <- cols do
      data_frame
      |> DataFrame.group_by(["sim_id"])
      |> DataFrame.frequencies([to_string(col)])
      |> DataFrame.rename([{col, "foo_id"}])
    end
    |> DataFrame.concat_rows()
    |> DataFrame.group_by(["sim_id", "foo_id"])
    |> DataFrame.summarise(cookies: sum(counts))
  end

  def do_rest(data_frame) do
    outcomes = ["overbaked", "underbaked"]

    data_frame
    |> DataFrame.group_by([:sim_id, :foo_id])
    |> DataFrame.summarise(
      cakes:
        sum(
          # 😧
          cond do
            result == "overbaked" or ingredients >= 2 ->
              0

              # like 5 more of these
          end
        ),
      biscuits: sum(if result in ^outcomes, do: 1, else: 0),
      pies: sum(if result == "overbaked", do: 1, else: 0)
    )
    |> DataFrame.mutate(hot_treats: biscuits * 4 + pies * 3)
  end
end

Name                ips        average  deviation         median         99th %
3_join           224.09        4.46 ms     ±6.97%        4.44 ms        5.39 ms
2_rest             1.22      818.17 ms     ±1.26%      811.82 ms      838.34 ms
1_cookies          1.13      882.87 ms     ±7.32%      863.68 ms      991.60 ms

Comparison: 
3_join           224.09
2_rest             1.22 - 183.35x slower +813.71 ms
1_cookies          1.13 - 197.85x slower +878.41 ms

Memory usage statistics:

Name         Memory usage
3_join            8.64 KB
2_rest           31.05 KB - 3.59x memory usage +22.41 KB
1_cookies        83.63 KB - 9.68x memory usage +74.98 KB

Python code & results (only for a subset of the full Elixir example)

import polars as pl
df = pl.read_parquet("/Users/billy/projects/elixir-explorer/explorer/motch.parquet")

def do_rest():
    return (
        df.group_by([pl.col("sim_id"), pl.col("foo_id")])
        .agg(
            cakes=pl.when(pl.Expr.or_((pl.col("result") == "overbaked"), pl.col("ingredients") > 2)).then(0).sum(),
            biscuits=pl.when(pl.col("result").is_in(["overbaked", "underbaked"])).then(1).otherwise(0).sum(),
            pies=pl.when(pl.col("result") == "overbaked").then(1).otherwise(0).sum(),
        )
    )

timeit.timeit("do_rest()", globals=locals(), number=100) / 100
#=> 0.03980570465093478

Summary

• Elixir: 800ms
• Python: 40ms

Not sure what's up yet. I'm guessing the query plans will reveal some bad defaults.

[mhanberg]

Anecdotally I'll say I changed the usage of group_by_stable to group_by (in the rust bindings) and saw a 20% improvement in my benchmarking.

Might be something else tho, as the timing here seems more extreme.

[mhanberg]

I was looking into this and ran the sample above on my M3 Pro macbook and got the following results

Operating System: macOS
CPU Information: Apple M3 Pro
Number of Available Cores: 12
Available memory: 36 GB
Elixir 1.16.2
Erlang 26.2.5
JIT enabled: true

Benchmark suite executing with the following configuration:
warmup: 2 s
time: 10 s
memory time: 2 s
reduction time: 0 ns
parallel: 1
inputs: none specified
Estimated total run time: 42 s

Benchmarking 1_cookies ...
Benchmarking 2_rest ...
Benchmarking 3_join ...
Calculating statistics...
Formatting results...

Name                ips        average  deviation         median         99th %
3_join          2388.72        0.42 ms    ±12.87%        0.41 ms        0.58 ms
2_rest            70.85       14.12 ms     ±4.64%       14.00 ms       16.56 ms
1_cookies         53.21       18.79 ms    ±10.62%       18.30 ms       28.15 ms

Comparison:
3_join          2388.72
2_rest            70.85 - 33.72x slower +13.70 ms
1_cookies         53.21 - 44.89x slower +18.38 ms

Memory usage statistics:

Name         Memory usage
3_join            4.45 KB
2_rest           24.10 KB - 5.41x memory usage +19.65 KB
1_cookies        55.77 KB - 12.52x memory usage +51.31 KB

[mhanberg]

Also, I'm am curious by this line in the python script

timeit.timeit("do_rest()", globals=locals(), number=100) / 100
#=> 0.03980570465093478

From what I've seen, timeit returns the time in seconds, so if you divide that by 100, you actually get it in hectoseconds. (if you want milliseconds, multiply by 1000).

So the result you got was actually about 3.9 seconds.

When i ran the same (slightly modified code), i get 1.56 seconds.

import polars as pl
import timeit
df = pl.read_parquet("/Users/m.hanberg/src/simplebet/pharmacy/clients/lib_sgp/motch.parquet")

def do_rest():
    return (
        df.group_by([pl.col("sim_id"), pl.col("foo_id")])
        .agg(
            cakes=pl.when(pl.Expr.or_((pl.col("result") == "overbaked"), pl.col("ingredients") > 2)).then(0).sum(),
            biscuits=pl.when(pl.col("result").is_in(["overbaked", "underbaked"])).then(1).otherwise(0).sum(),
            pies=pl.when(pl.col("result") == "overbaked").then(1).otherwise(0).sum(),
        )
    )

value = timeit.timeit("do_rest()", globals=locals(), number=100)
print(value)

and run it with uv uv --with polars test.py.

[josevalim]

I don't know if this came up before but I wonder if a lazy dataframe would improve considerably your results.

[mhanberg]

For these examples, I made sure they were all fully lazy, but the results came up the same as before

[billylanchantin]

Also, I'm am curious by this line in the python script

timeit.timeit("do_rest()", globals=locals(), number=100) / 100
#=> 0.03980570465093478

The docs say that number tells timeit.timeit how many times to execute the main statement:

• https://docs.python.org/3/library/timeit.html#timeit.Timer.timeit

So when I divide by 100, that's not a unit conversion but rather it gives the average execution time over 100 attempts.

[billylanchantin]

@mhanberg If I'm reading this right, looks like you got similar performance!

• Your Elixir: 14.12ms
• Your Python: 15.60ms

Did you still have your group_by set to unstable? If so that's telling.

[billylanchantin]

TODO: Revisit this. Now that #1102 has been merged, I suspect this is no longer an issue.