#216 - fix incorrect sorting

business-science · Nov 5, 2023 · a134a07 · a134a07
1 parent 00be13f
commit a134a07
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 44 deletions.
diff --git a/src/pytimetk/feature_engineering/expanding.py b/src/pytimetk/feature_engineering/expanding.py
@@ -522,29 +522,19 @@ def _augment_expanding_polars(
             expanding_exprs.append(expanding_expr)
             new_column_names.append(new_column_name)
 
-    # Convert Pandas DataFrame to Polars and ensure a consistent row order by resetting the index
-    df = pl.from_pandas(pandas_df.reset_index())
-
-    # Evaluate the accumulated expanding expressions and convert back to a Pandas DataFrame
-    if group_names:
-        df_new_columns = df \
-            .sort(*group_names, date_column) \
-            .group_by(group_names) \
-            .agg(expanding_exprs) \
-            .sort(*group_names) \
-            .explode(new_column_names)
+    # Select the columns
+    selected_columns = expanding_exprs
+
+    df = pl.DataFrame(pandas_df)
+    if isinstance(data, pd.core.groupby.generic.DataFrameGroupBy):
+        out_df = df.group_by(data.grouper.names, maintain_order=True).agg(selected_columns)
+        out_df = out_df.explode(out_df.columns[1:])
+        out_df = out_df.drop(data.grouper.names)
+    else: # a dataframe
+        out_df = df.select(selected_columns)
 
-        df = pl.concat([df, df_new_columns.drop(group_names)], how="horizontal") \
-                .sort('index') \
-                .drop('index') \
-                .to_pandas()
-    else:
-        df = df \
-            .sort(date_column) \
-            .with_columns(expanding_exprs) \
-            .sort('index') \
-            .drop('index') \
-            .to_pandas()
+    # Concatenate the DataFrames horizontally
+    df = pl.concat([df, out_df], how="horizontal").to_pandas()
 
     return reduce_memory_usage(df)
 

diff --git a/src/pytimetk/feature_engineering/rolling.py b/src/pytimetk/feature_engineering/rolling.py
@@ -514,30 +514,20 @@ def _augment_rolling_polars(
                 # Add constructed expressions and new column names to respective lists
                 rolling_exprs.append(rolling_expr)
                 new_column_names.append(new_column_name)
-
-    # Convert Pandas DataFrame to Polars and ensure a consistent row order by resetting the index
-    df = pl.from_pandas(pandas_df.reset_index())
 
-    # Evaluate the accumulated rolling expressions and convert back to a Pandas DataFrame
-    if group_names:
-        df_new_columns = df \
-            .sort(*group_names, date_column) \
-            .group_by(group_names) \
-            .agg(rolling_exprs) \
-            .sort(*group_names) \
-            .explode(new_column_names)
+    # Select the columns
+    selected_columns = rolling_exprs
+
+    df = pl.DataFrame(pandas_df)
+    if isinstance(data, pd.core.groupby.generic.DataFrameGroupBy):
+        out_df = df.group_by(data.grouper.names, maintain_order=True).agg(selected_columns)
+        out_df = out_df.explode(out_df.columns[1:])
+        out_df = out_df.drop(data.grouper.names)
+    else: # a dataframe
+        out_df = df.select(selected_columns)
 
-        df = pl.concat([df, df_new_columns.drop(group_names)], how="horizontal") \
-                .sort('index') \
-                .drop('index') \
-                .to_pandas()
-    else:
-        df = df \
-            .sort(date_column) \
-            .with_columns(rolling_exprs) \
-            .sort('index') \
-            .drop('index') \
-            .to_pandas()
+    # Concatenate the DataFrames horizontally
+    df = pl.concat([df, out_df], how="horizontal").to_pandas()
 
     return reduce_memory_usage(df)