feat(rust, python): in rolling aggregation functions, sort under-the-…

…hood for user if data is unsorted (with warning)

crates/polars-plan/src/dsl/mod.rs

@@ -1230,19 +1230,47 @@ impl Expr {
                 move |s| {
                     let mut by = s[1].clone();
                     by = by.rechunk();
-                    let s = &s[0];
+                    let series: Series;
 
                     polars_ensure!(
                         options.weights.is_none(),
                         ComputeError: "`weights` is not supported in 'rolling by' expression"
                     );
-                    let (by, tz) = match by.dtype() {
+                    let (mut by, tz) = match by.dtype() {
                         DataType::Datetime(tu, tz) => {
                             (by.cast(&DataType::Datetime(*tu, None))?, tz)
                         },
                         _ => (by.clone(), &None),
                     };
-                    ensure_sorted_arg(&by, expr_name)?;
+                    let sorting_indices;
+                    let original_indices;
+                    let by_flag = by.is_sorted_flag();
+                    if matches!(by_flag, IsSorted::Ascending) {
+                        series = s[0].clone();
+                        original_indices = None;
+                    }
+                    else if matches!(by_flag, IsSorted::Descending) {
+                        series = s[0].reverse();
+                        by = by.reverse();
+                        original_indices = None;
+                    }
+                    else {
+                        if options.warn_if_unsorted {
+                            eprintln!(
+                                "PolarsPerformanceWarning: Series is not known to be sorted by `by` column, \
+                                so Polars will temporarily sort it for you.\n\
+                                You can silence this warning by:\n\
+                                - passing `warn_if_unsorted=False`;\n\
+                                - sorting your data by your `by` column beforehand;\n\
+                                - setting `.set_sorted()` if you already know your data is sorted\n\
+                                before passing it to the rolling aggregation function"
+                            );
+                        }
+                        sorting_indices = by.arg_sort(Default::default());
+                        by = by.take(&sorting_indices)?;
+                        series = s[0].take(&sorting_indices)?;
+                        original_indices = UInt32Chunked::from_iter_values("", 0..s[0].len() as u32).into_series().take(&sorting_indices).ok();
+                    };
                     let by = by.datetime().unwrap();
                     let by_values = by.cont_slice().map_err(|_| {
                         polars_err!(
@@ -1264,7 +1292,13 @@ impl Expr {
                         fn_params: options.fn_params.clone(),
                     };
 
-                    rolling_fn(s, options).map(Some)
+                    if matches!(by_flag, IsSorted::Ascending) {
+                        rolling_fn(&series, options).map(Some)
+                    } else if matches!(by_flag, IsSorted::Descending) {
+                        Ok(rolling_fn(&series, options)?.reverse()).map(Some)
+                    } else {
+                        rolling_fn(&series, options)?.take(&original_indices.unwrap().arg_sort(Default::default())).map(Some)
+                    }
                 },
                 &[col(by)],
                 output_type,

crates/polars-time/src/chunkedarray/rolling_window/mod.rs

@@ -38,6 +38,8 @@ pub struct RollingOptions {
     pub closed_window: Option<ClosedWindow>,
     /// Optional parameters for the rolling function
     pub fn_params: DynArgs,
+    /// Warn if data is not known to be sorted by `by` column (if passed)
+    pub warn_if_unsorted: bool,
 }
 
 #[cfg(feature = "rolling_window")]
@@ -51,6 +53,7 @@ impl Default for RollingOptions {
             by: None,
             closed_window: None,
             fn_params: None,
+            warn_if_unsorted: true,
         }
     }
 }

py-polars/polars/expr/expr.py

@@ -5279,6 +5279,7 @@ def rolling_min(
         center: bool = False,
         by: str | None = None,
         closed: ClosedInterval = "left",
+        warn_if_unsorted: bool = True,
     ) -> Self:
         """
         Apply a rolling min (moving min) over the values in this array.
@@ -5350,6 +5351,8 @@ def rolling_min(
         closed : {'left', 'right', 'both', 'none'}
             Define which sides of the temporal interval are closed (inclusive); only
             applicable if `by` has been set.
+        warn_if_unsorted
+            Warn if data is not known to be sorted by `by` column (if passed).
 
         Warnings
         --------
@@ -5475,7 +5478,7 @@ def rolling_min(
         )
         return self._from_pyexpr(
             self._pyexpr.rolling_min(
-                window_size, weights, min_periods, center, by, closed
+                window_size, weights, min_periods, center, by, closed, warn_if_unsorted
             )
         )
 
@@ -5489,6 +5492,7 @@ def rolling_max(
         center: bool = False,
         by: str | None = None,
         closed: ClosedInterval = "left",
+        warn_if_unsorted: bool = True,
     ) -> Self:
         """
         Apply a rolling max (moving max) over the values in this array.
@@ -5556,6 +5560,8 @@ def rolling_max(
         closed : {'left', 'right', 'both', 'none'}
             Define which sides of the temporal interval are closed (inclusive); only
             applicable if `by` has been set.
+        warn_if_unsorted
+            Warn if data is not known to be sorted by `by` column (if passed).
 
         Warnings
         --------
@@ -5708,7 +5714,7 @@ def rolling_max(
         )
         return self._from_pyexpr(
             self._pyexpr.rolling_max(
-                window_size, weights, min_periods, center, by, closed
+                window_size, weights, min_periods, center, by, closed, warn_if_unsorted
             )
         )
 
@@ -5722,6 +5728,7 @@ def rolling_mean(
         center: bool = False,
         by: str | None = None,
         closed: ClosedInterval = "left",
+        warn_if_unsorted: bool = True,
     ) -> Self:
         """
         Apply a rolling mean (moving mean) over the values in this array.
@@ -5793,6 +5800,8 @@ def rolling_mean(
         closed : {'left', 'right', 'both', 'none'}
             Define which sides of the temporal interval are closed (inclusive); only
             applicable if `by` has been set.
+        warn_if_unsorted
+            Warn if data is not known to be sorted by `by` column (if passed).
 
         Warnings
         --------
@@ -5945,7 +5954,13 @@ def rolling_mean(
         )
         return self._from_pyexpr(
             self._pyexpr.rolling_mean(
-                window_size, weights, min_periods, center, by, closed
+                window_size,
+                weights,
+                min_periods,
+                center,
+                by,
+                closed,
+                warn_if_unsorted,
             )
         )
 
@@ -5959,6 +5974,7 @@ def rolling_sum(
         center: bool = False,
         by: str | None = None,
         closed: ClosedInterval = "left",
+        warn_if_unsorted: bool = True,
     ) -> Self:
         """
         Apply a rolling sum (moving sum) over the values in this array.
@@ -6026,6 +6042,8 @@ def rolling_sum(
         closed : {'left', 'right', 'both', 'none'}
             Define which sides of the temporal interval are closed (inclusive); only
             applicable if `by` has been set.
+        warn_if_unsorted
+            Warn if data is not known to be sorted by `by` column (if passed).
 
         Warnings
         --------
@@ -6178,7 +6196,7 @@ def rolling_sum(
         )
         return self._from_pyexpr(
             self._pyexpr.rolling_sum(
-                window_size, weights, min_periods, center, by, closed
+                window_size, weights, min_periods, center, by, closed, warn_if_unsorted
             )
         )
 
@@ -6193,6 +6211,7 @@ def rolling_std(
         by: str | None = None,
         closed: ClosedInterval = "left",
         ddof: int = 1,
+        warn_if_unsorted: bool = True,
     ) -> Self:
         """
         Compute a rolling standard deviation.
@@ -6262,6 +6281,8 @@ def rolling_std(
             applicable if `by` has been set.
         ddof
             "Delta Degrees of Freedom": The divisor for a length N window is N - ddof
+        warn_if_unsorted
+            Warn if data is not known to be sorted by `by` column (if passed).
 
         Warnings
         --------
@@ -6414,7 +6435,14 @@ def rolling_std(
         )
         return self._from_pyexpr(
             self._pyexpr.rolling_std(
-                window_size, weights, min_periods, center, by, closed, ddof
+                window_size,
+                weights,
+                min_periods,
+                center,
+                by,
+                closed,
+                ddof,
+                warn_if_unsorted,
             )
         )
 
@@ -6429,6 +6457,7 @@ def rolling_var(
         by: str | None = None,
         closed: ClosedInterval = "left",
         ddof: int = 1,
+        warn_if_unsorted: bool = True,
     ) -> Self:
         """
         Compute a rolling variance.
@@ -6498,6 +6527,8 @@ def rolling_var(
             applicable if `by` has been set.
         ddof
             "Delta Degrees of Freedom": The divisor for a length N window is N - ddof
+        warn_if_unsorted
+            Warn if data is not known to be sorted by `by` column (if passed).
 
         Warnings
         --------
@@ -6657,6 +6688,7 @@ def rolling_var(
                 by,
                 closed,
                 ddof,
+                warn_if_unsorted,
             )
         )
 
@@ -6670,6 +6702,7 @@ def rolling_median(
         center: bool = False,
         by: str | None = None,
         closed: ClosedInterval = "left",
+        warn_if_unsorted: bool = True,
     ) -> Self:
         """
         Compute a rolling median.
@@ -6737,6 +6770,8 @@ def rolling_median(
         closed : {'left', 'right', 'both', 'none'}
             Define which sides of the temporal interval are closed (inclusive); only
             applicable if `by` has been set.
+        warn_if_unsorted
+            Warn if data is not known to be sorted by `by` column (if passed).
 
         Warnings
         --------
@@ -6815,7 +6850,7 @@ def rolling_median(
         )
         return self._from_pyexpr(
             self._pyexpr.rolling_median(
-                window_size, weights, min_periods, center, by, closed
+                window_size, weights, min_periods, center, by, closed, warn_if_unsorted
             )
         )
 
@@ -6831,6 +6866,7 @@ def rolling_quantile(
         center: bool = False,
         by: str | None = None,
         closed: ClosedInterval = "left",
+        warn_if_unsorted: bool = True,
     ) -> Self:
         """
         Compute a rolling quantile.
@@ -6902,6 +6938,8 @@ def rolling_quantile(
         closed : {'left', 'right', 'both', 'none'}
             Define which sides of the temporal interval are closed (inclusive); only
             applicable if `by` has been set.
+        warn_if_unsorted
+            Warn if data is not known to be sorted by `by` column (if passed).
 
         Warnings
         --------
@@ -7016,6 +7054,7 @@ def rolling_quantile(
                 center,
                 by,
                 closed,
+                warn_if_unsorted,
             )
         )