perf: Provide a general fast path for arg_sort_multiple.

crates/polars-core/src/frame/column/mod.rs

@@ -894,8 +894,263 @@ impl Column {
     }
 
     pub fn arg_sort(&self, options: SortOptions) -> IdxCa {
-        // @scalar-opt
-        self.as_materialized_series().arg_sort(options)
+        if self.is_empty() {
+            return IdxCa::from_vec(self.name().clone(), Vec::new());
+        }
+
+        if self.null_count() == self.len() {
+            // We might need to maintain order so just respect the descending parameter.
+            let values = if options.descending {
+                (0..self.len() as IdxSize).rev().collect()
+            } else {
+                (0..self.len() as IdxSize).collect()
+            };
+
+            return IdxCa::from_vec(self.name().clone(), values);
+        }
+
+        let is_sorted = Some(self.is_sorted_flag());
+        let Some(is_sorted) = is_sorted.filter(|v| !matches!(v, IsSorted::Not)) else {
+            return self.as_materialized_series().arg_sort(options);
+        };
+
+        // Fast path: the data is sorted.
+        let is_sorted_dsc = matches!(is_sorted, IsSorted::Descending);
+        let invert = options.descending != is_sorted_dsc;
+
+        let mut values = Vec::with_capacity(self.len());
+
+        #[inline(never)]
+        fn extend(
+            start: IdxSize,
+            end: IdxSize,
+            slf: &Column,
+            values: &mut Vec<IdxSize>,
+            is_only_nulls: bool,
+            invert: bool,
+            maintain_order: bool,
+        ) {
+            debug_assert!(start <= end);
+            debug_assert!(start as usize <= slf.len());
+            debug_assert!(end as usize <= slf.len());
+
+            if !invert || is_only_nulls {
+                values.extend(start..end);
+                return;
+            }
+
+            // If we don't have to maintain order but we have to invert. Just flip it around.
+            if !maintain_order {
+                values.extend((start..end).rev());
+                return;
+            }
+
+            // If we want to maintain order but we also needs to invert, we need to invert
+            // per group of items.
+            //
+            // @NOTE: Since the column is sorted, arg_unique can also take a fast path and
+            // just do a single traversal.
+            let arg_unique = slf
+                .slice(start as i64, (end - start) as usize)
+                .arg_unique()
+                .unwrap();
+
+            assert!(!arg_unique.has_nulls());
+
+            let num_unique = arg_unique.len();
+
+            // Fast path: all items are unique.
+            if num_unique == (end - start) as usize {
+                values.extend((start..end).rev());
+                return;
+            }
+
+            if num_unique == 1 {
+                values.extend(start..end);
+                return;
+            }
+
+            let mut prev_idx = end - start;
+            for chunk in arg_unique.downcast_iter() {
+                for &idx in chunk.values().as_slice().iter().rev() {
+                    values.extend(start + idx..start + prev_idx);
+                    prev_idx = idx;
+                }
+            }
+        }
+        macro_rules! extend {
+            ($start:expr, $end:expr) => {
+                extend!($start, $end, is_only_nulls = false);
+            };
+            ($start:expr, $end:expr, is_only_nulls = $is_only_nulls:expr) => {
+                extend(
+                    $start,
+                    $end,
+                    self,
+                    &mut values,
+                    $is_only_nulls,
+                    invert,
+                    options.maintain_order,
+                );
+            };
+        }
+
+        let length = self.len() as IdxSize;
+        let null_count = self.null_count() as IdxSize;
+
+        if null_count == 0 {
+            extend!(0, length);
+        } else {
+            let has_nulls_last = self.get(self.len() - 1).unwrap().is_null();
+            match (options.nulls_last, has_nulls_last) {
+                (true, true) => {
+                    // Current: Nulls last, Wanted: Nulls last
+                    extend!(0, length - null_count);
+                    extend!(length - null_count, length, is_only_nulls = true);
+                },
+                (true, false) => {
+                    // Current: Nulls first, Wanted: Nulls last
+                    extend!(null_count, length);
+                    extend!(0, null_count, is_only_nulls = true);
+                },
+                (false, true) => {
+                    // Current: Nulls last, Wanted: Nulls first
+                    extend!(length - null_count, length, is_only_nulls = true);
+                    extend!(0, length - null_count);
+                },
+                (false, false) => {
+                    // Current: Nulls first, Wanted: Nulls first
+                    extend!(0, null_count, is_only_nulls = true);
+                    extend!(null_count, length);
+                },
+            }
+        }
+
+        // @NOTE: This can theoretically be pushed into the previous operation but it is really
+        // worth it... probably not...
+        if let Some((limit, limit_dsc)) = options.limit {
+            let limit = limit.min(length);
+
+            if limit_dsc {
+                values = values.drain((length - limit) as usize..).collect();
+            } else {
+                values.truncate(limit as usize);
+            }
+        }
+
+        IdxCa::from_vec(self.name().clone(), values)
+    }
+
+    pub fn arg_sort_multiple(
+        &self,
+        by: &[Column],
+        options: &SortMultipleOptions,
+    ) -> PolarsResult<IdxCa> {
+        /// Determine whether we can skip a column in the sorting process.
+        fn is_column_skippable(
+            c: &Column,
+            descending: bool,
+            nulls_last: bool,
+            maintain_order: bool,
+        ) -> bool {
+            if c.as_scalar_column().is_some() {
+                return true;
+            }
+
+            if maintain_order {
+                return false;
+            }
+
+            let is_sorted = c.is_sorted_flag();
+            if matches!(is_sorted, IsSorted::Not) {
+                return false;
+            }
+
+            let is_sorted_dsc = matches!(is_sorted, IsSorted::Descending);
+            if descending != is_sorted_dsc {
+                // @TODO: This can probably be handled smarter by instead inserting a column with
+                // inverse indices.
+                return false;
+            }
+
+            if !c.has_nulls() {
+                return true;
+            }
+
+            let has_nulls_last = c.get(c.len() - 1).unwrap().is_null();
+            if nulls_last == has_nulls_last {
+                return true;
+            }
+
+            false
+        }
+
+        let mut descending = Vec::with_capacity(by.len() + 1);
+        let mut nulls_last = Vec::with_capacity(by.len() + 1);
+        let mut non_skippable_by = Vec::with_capacity(by.len() + 1);
+
+        for (i, c) in std::iter::once(self).chain(by.iter()).enumerate() {
+            let dsc = options
+                .descending
+                .get(i)
+                .cloned()
+                .unwrap_or(options.descending[0]);
+            let nl = options
+                .nulls_last
+                .get(i)
+                .cloned()
+                .unwrap_or(options.nulls_last[0]);
+
+            if is_column_skippable(c, dsc, nl, options.maintain_order) {
+                continue;
+            }
+
+            descending.push(dsc);
+            nulls_last.push(nl);
+            non_skippable_by.push(c.clone());
+        }
+
+        if non_skippable_by.is_empty() {
+            return Ok(IdxCa::new_vec(
+                self.name().clone(),
+                (0..self.len() as IdxSize).collect(),
+            ));
+        }
+
+        if non_skippable_by.len() == 1 {
+            return Ok(non_skippable_by[0]
+                .arg_sort(SortOptions {
+                    descending: descending[0],
+                    nulls_last: nulls_last[0],
+                    multithreaded: options.multithreaded,
+                    maintain_order: options.maintain_order,
+                    limit: options.limit,
+                })
+                .with_name(self.name().clone()));
+        }
+
+        let by = &non_skippable_by[1..];
+
+        Ok(non_skippable_by[0]
+            .as_materialized_series()
+            .arg_sort_multiple(
+                by,
+                &SortMultipleOptions {
+                    descending,
+                    nulls_last,
+                    multithreaded: options.multithreaded,
+                    maintain_order: options.maintain_order,
+                    limit: options.limit,
+                },
+            )?
+            .with_name(self.name().clone()))
+    }
+
+    pub fn arg_unique(&self) -> PolarsResult<IdxCa> {
+        match self {
+            Column::Scalar(s) => Ok(IdxCa::new_vec(s.name().clone(), vec![0])),
+            _ => self.as_materialized_series().arg_unique(),
+        }
     }
 
     pub fn bit_repr(&self) -> Option<BitRepr> {
@@ -986,8 +1241,11 @@ impl Column {
     }
 
     pub fn is_sorted_flag(&self) -> IsSorted {
-        // @scalar-opt
-        self.as_materialized_series().is_sorted_flag()
+        match self {
+            Column::Series(s) => s.is_sorted_flag(),
+            Column::Partitioned(s) => s.partitions().is_sorted_flag(),
+            Column::Scalar(_) => IsSorted::Ascending,
+        }
     }
 
     pub fn unique(&self) -> PolarsResult<Column> {

crates/polars-core/src/frame/mod.rs

@@ -2143,7 +2143,6 @@ impl DataFrame {
                 } else {
                     let (first, other) = prepare_arg_sort(by_column, &mut sort_options)?;
                     first
-                        .as_materialized_series()
                         .arg_sort_multiple(&other, &sort_options)?
                 }
             },

crates/polars-expr/src/expressions/sortby.rs

@@ -163,10 +163,7 @@ fn sort_by_groups_multiple_by(
                 limit: None,
             };
 
-            let sorted_idx = groups[0]
-                .as_materialized_series()
-                .arg_sort_multiple(&groups[1..], &options)
-                .unwrap();
+            let sorted_idx = groups[0].arg_sort_multiple(&groups[1..], &options).unwrap();
             map_sorted_indices_to_group_idx(&sorted_idx, idx)
         },
         GroupsIndicator::Slice([first, len]) => {
@@ -183,10 +180,7 @@ fn sort_by_groups_multiple_by(
                 maintain_order,
                 limit: None,
             };
-            let sorted_idx = groups[0]
-                .as_materialized_series()
-                .arg_sort_multiple(&groups[1..], &options)
-                .unwrap();
+            let sorted_idx = groups[0].arg_sort_multiple(&groups[1..], &options).unwrap();
             map_sorted_indices_to_group_slice(&sorted_idx, first)
         },
     };
@@ -246,9 +240,7 @@ impl PhysicalExpr for SortByExpr {
                     );
                 }
 
-                s_sort_by[0]
-                    .as_materialized_series()
-                    .arg_sort_multiple(&s_sort_by[1..], &options)
+                s_sort_by[0].arg_sort_multiple(&s_sort_by[1..], &options)
             };
             POOL.install(|| rayon::join(series_f, sorted_idx_f))
         };

py-polars/tests/unit/dataframe/test_df.py

@@ -1014,7 +1014,7 @@ def test_multiple_column_sort() -> None:
         pl.DataFrame({"a": [3, 2, 1], "b": ["b", "a", "a"]}),
     )
     assert_frame_equal(
-        df.sort("b", descending=True),
+        df.sort("b", descending=True, maintain_order=True),
         pl.DataFrame({"a": [3, 1, 2], "b": ["b", "a", "a"]}),
     )
     assert_frame_equal(

py-polars/tests/unit/datatypes/test_categorical.py

@@ -825,6 +825,8 @@ def test_cat_preserve_lexical_ordering_on_concat() -> None:
     assert df2["x"].dtype == dtype
 
 
+# TODO: Bug see: https://github.com/pola-rs/polars/issues/20440
+@pytest.mark.may_fail_auto_streaming
 def test_cat_append_lexical_sorted_flag() -> None:
     df = pl.DataFrame({"x": [0, 1, 1], "y": ["B", "B", "A"]}).with_columns(
         pl.col("y").cast(pl.Categorical(ordering="lexical"))

py-polars/tests/unit/operations/namespaces/test_categorical.py

@@ -29,6 +29,8 @@ def test_global_and_local(
         yield
 
 
+# @TODO: Bug, see https://github.com/pola-rs/polars/issues/20440
+@pytest.mark.may_fail_auto_streaming
 def test_categorical_lexical_sort() -> None:
     df = pl.DataFrame(
         {"cats": ["z", "z", "k", "a", "b"], "vals": [3, 1, 2, 2, 3]}

py-polars/tests/unit/operations/test_interpolate_by.py

@@ -105,17 +105,15 @@ def test_interpolate_by_leading_nulls() -> None:
         }
     )
     result = df.select(pl.col("values").interpolate_by("times"))
-    expected = pl.DataFrame(
-        {"values": [None, None, None, 1.0, 1.7999999999999998, 4.6, 5.0]}
-    )
+    expected = pl.DataFrame({"values": [None, None, None, 1.0, 1.8, 4.6, 5.0]})
     assert_frame_equal(result, expected)
     result = (
-        df.sort("times", descending=True)
+        df.sort("times", maintain_order=True, descending=True)
         .with_columns(pl.col("values").interpolate_by("times"))
-        .sort("times")
+        .sort("times", maintain_order=True)
         .drop("times")
     )
-    assert_frame_equal(result, expected)
+    assert_frame_equal(result, expected, check_exact=False)
 
 
 @pytest.mark.parametrize("dataset", ["floats", "dates"])

py-polars/tests/unit/operations/test_join.py

@@ -285,11 +285,18 @@ def test_join_on_cast() -> None:
 
     df_b = pl.DataFrame({"a": [-2, -3, 3, 10]})
 
-    assert df_a.join(df_b, on=pl.col("a").cast(pl.Int64)).to_dict(as_series=False) == {
-        "index": [1, 2, 3, 5],
-        "a": [-2, 3, 3, 10],
-        "a_right": [-2, 3, 3, 10],
-    }
+    assert_frame_equal(
+        df_a.join(df_b, on=pl.col("a").cast(pl.Int64)),
+        pl.DataFrame(
+            {
+                "index": [1, 2, 3, 5],
+                "a": [-2, 3, 3, 10],
+                "a_right": [-2, 3, 3, 10],
+            }
+        ),
+        check_row_order=False,
+        check_dtypes=False,
+    )
     assert df_a.lazy().join(
         df_b.lazy(), on=pl.col("a").cast(pl.Int64)
     ).collect().to_dict(as_series=False) == {