From 04f1d377c9b5241c8f6265242b25bae0f2a5f15d Mon Sep 17 00:00:00 2001
From: Rian Koja <riankoja@gmail.com>
Date: Tue, 7 Oct 2025 17:28:08 -0300
Subject: [PATCH] feat: Incorporated Saha Sourav's feedback

---
 00_tools.py     |  4 +++-
 02_benchmark.py | 16 ++++++++++++++--
 03_polars.py    | 22 +++++++++++++++++++---
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/00_tools.py b/00_tools.py
index 40f2d2e..71d317b 100644
--- a/00_tools.py
+++ b/00_tools.py
@@ -21,7 +21,9 @@ def time_operation(
     start_time = time.perf_counter()
     result = func(*args, **kwargs)
     # Force evaluation for lazy operations
-    if hasattr(result, "compute"):
+    if hasattr(result, "_evaluate"):
+        result = result._evaluate()
+    elif hasattr(result, "compute"):
         result = result.compute()
     elif hasattr(result, "values"):
         _ = result.values  # Access values to force computation
diff --git a/02_benchmark.py b/02_benchmark.py
index 63c7909..3a9ef01 100644
--- a/02_benchmark.py
+++ b/02_benchmark.py
@@ -123,10 +123,10 @@ def groupby_aggregation_operation():
         )
     )
 
-    # Window functions
+    # groupby-dense-rank
     results.append(
         time_operation(
-            "window_functions",
+            "groupby-dense-rank",
             df_lib,
             lambda: orders.assign(
                 running_total=orders.groupby("customer_id")["total_amount"].cumsum(),
@@ -135,6 +135,18 @@ def groupby_aggregation_operation():
         )
     )
 
+    # groupby-first-rank
+    results.append(
+        time_operation(
+            "groupby-first-rank",
+            df_lib,
+            lambda: orders.assign(
+                running_total=orders.groupby("customer_id")["total_amount"].cumsum(),
+                rank=orders.groupby("customer_id")["total_amount"].rank(method="first"),
+            ),
+        )
+    )
+
     # String operations
     results.append(
         time_operation(
diff --git a/03_polars.py b/03_polars.py
index 9863d58..b0168c6 100644
--- a/03_polars.py
+++ b/03_polars.py
@@ -137,8 +137,8 @@ def four_table_join_polars():
 
     results.append(time_operation("four_table_join", pl, four_table_join_polars))
 
-    # Window functions
-    def window_functions_polars():
+    # groupby-dense-rank
+    def groupby_dense_rank_polars():
         # Use Polars native window functions.
         # Cast rank to float to match pandas output dtype.
         result = orders.with_columns(
@@ -151,7 +151,23 @@ def window_functions_polars():
         )
         return result
 
-    results.append(time_operation("window_functions", pl, window_functions_polars))
+    results.append(time_operation("groupby-dense-rank", pl, groupby_dense_rank_polars))
+
+    # groupby-first-rank
+    def groupby_first_rank_polars():
+        # Use Polars native window functions with first ranking method.
+        # Cast rank to float to match pandas output dtype.
+        result = orders.with_columns(
+            pl.col("total_amount").cum_sum().over("customer_id").alias("running_total"),
+            pl.col("total_amount")
+            .rank(method="ordinal")
+            .over("customer_id")
+            .cast(pl.Float64)
+            .alias("rank"),
+        )
+        return result
+
+    results.append(time_operation("groupby-first-rank", pl, groupby_first_rank_polars))
 
     # String operations
     def string_operations_polars():