feat: Incorporated Saha Sourav's feedback

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

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(

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():