Add groupBy benchmark.

Author: mchav

benchmark/Main.hs

@@ -7,6 +7,7 @@ import qualified Data.Vector.Unboxed.Mutable as VUM
 
 import Control.Monad (replicateM)
 import Criterion.Main
+import DataFrame ((|>))
 import Data.Time
 import System.Process
 import System.Random.Stateful
@@ -55,10 +56,31 @@ pandas = do
   output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/pandas/pandas_benchmark.py"] ""
   putStrLn output
 
+groupByHaskell :: IO ()
+groupByHaskell = do
+  df <- D.readCsv "./data/housing.csv"
+  print $ df |> D.groupBy ["ocean_proximity"]
+             |> D.aggregate [("median_house_value", D.Minimum), ("median_house_value", D.Maximum)]
+             |> D.select ["ocean_proximity", "Minimum_median_house_value", "Maximum_median_house_value"]
+
+groupByPolars :: IO ()
+groupByPolars = do
+  output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/polars/group_by.py"] ""
+  putStrLn output
+
+groupByPandas :: IO ()
+groupByPandas = do
+  output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/pandas/group_by.py"] ""
+  putStrLn output
+
 main = do
   defaultMain [
-    bgroup "stats" [ bench  "haskell" $ nfIO haskell
-                   , bench  "polars"  $ nfIO polars
-                   , bench  "pandas"  $ nfIO pandas
+    bgroup "stats" [ 
+                   , bench  "simpleStatsHaskell" $ nfIO haskell
+                   , bench  "simpleStatsPandas" $ nfIO pandas
+                   , bench  "simpleStatsPolars" $ nfIO polars
+                   , bench  "groupByHaskell" $ nfIO groupByHaskell
+                   , bench  "groupByPolars"  $ nfIO groupByPolars
+                   , bench  "groupByPandas"  $ nfIO groupByPandas
                    ]
     ]

benchmark/pandas/group_by.py

@@ -0,0 +1,15 @@
+import pandas as pd
+
+df = pd.read_csv("./data/housing.csv")
+
+# Group, aggregate, and rename in one shot
+agg_df = (
+    df
+    .groupby("ocean_proximity")["median_house_value"]
+    .agg(
+        Minimum_median_house_value="min",
+        Maximum_median_house_value="max"
+    )
+)
+
+print(agg_df)

benchmark/polars/group_by.py

@@ -0,0 +1,18 @@
+import datetime
+import numpy as np
+import polars as pl
+
+# ------------------------------------------------------------------------------
+
+df = pl.read_csv("./data/housing.csv")
+
+result = (
+    df
+    .group_by("ocean_proximity")
+    .agg([
+        pl.col("median_house_value").min().alias("Minimum_median_house_value"),
+        pl.col("median_house_value").max().alias("Maximum_median_house_value")
+    ])
+)
+
+print(result)