Merge branch 'main' into arrow-pushdown-pyo

Author: tylerriccio33

.github/workflows/lint-rust.yml

@@ -103,3 +103,15 @@ jobs:
           --features object
           -p polars-core
 #          -p polars-arrow
+
+  deny:
+    if: github.ref_name != 'main'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Install cargo-deny
+        uses: taiki-e/install-action@cargo-deny
+
+      - name: Run cargo deny
+        run: cargo deny check

Makefile

@@ -160,6 +160,10 @@ fmt: update-cargo-env ## Run autoformatting and linting
 	dprint fmt
 	$(VENV_BIN)/typos
 
+.PHONY: deny
+deny:
+	cargo deny check
+
 .PHONY: fix
 fix: update-cargo-env
 	cargo clippy --workspace --all-targets --all-features --fix

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

@@ -425,11 +425,23 @@ impl PhysicalExpr for AggregationExpr {
                     AggregatedScalar(agg_s.with_name(keep_name))
                 },
                 GroupByMethod::Implode { maintain_order: _ } => {
-                    AggregatedScalar(match ac.agg_state() {
+                    let col = match ac.agg_state() {
                         AggState::LiteralScalar(_) => unreachable!(), // handled above
                         AggState::AggregatedScalar(c) => c.as_list().into_column(),
-                        AggState::NotAggregated(_) | AggState::AggregatedList(_) => ac.aggregated(),
-                    })
+                        AggState::AggregatedList(c) => c.clone(),
+                        AggState::NotAggregated(_) => ac.aggregated(),
+                    };
+                    // TODO: Introduce `UpdateGroups::WithUnitLen` as a new lazy `groups()` method
+                    // and move the groups constructor there. Then, set `UpdateGroups::WithUnitLen` to
+                    // all AggregationExprs.
+                    let groups = Cow::Owned({
+                        let groups = (0..col.len() as IdxSize).map(|i| [i, 1]).collect();
+                        GroupsType::new_slice(groups, false, true).into_sliceable()
+                    });
+                    return Ok(AggregationContext::from_agg_state(
+                        AggregatedScalar(col),
+                        groups,
+                    ));
                 },
                 GroupByMethod::Groups => {
                     let mut column: ListChunked = ac.groups().as_list_chunked();

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

@@ -1942,7 +1942,12 @@ impl LazyFrame {
     }
 
     #[cfg(feature = "merge_sorted")]
-    pub fn merge_sorted<S>(self, other: LazyFrame, key: S) -> PolarsResult<LazyFrame>
+    pub fn merge_sorted<S>(
+        self,
+        other: LazyFrame,
+        key: S,
+        maintain_order: bool,
+    ) -> PolarsResult<LazyFrame>
     where
         S: Into<PlSmallStr>,
     {
@@ -1952,6 +1957,7 @@ impl LazyFrame {
             input_left: Arc::new(self.logical_plan),
             input_right: Arc::new(other.logical_plan),
             key,
+            maintain_order,
         };
         Ok(LazyFrame::from_logical_plan(lp, self.opt_state))
     }

crates/polars-mem-engine/src/executors/sort.rs

@@ -10,8 +10,13 @@ pub(crate) struct SortExec {
 }
 
 impl SortExec {
-    fn execute_impl(&mut self, state: &ExecutionState, df: DataFrame) -> PolarsResult<DataFrame> {
+    fn execute_impl(
+        &mut self,
+        state: &ExecutionState,
+        mut df: DataFrame,
+    ) -> PolarsResult<DataFrame> {
         state.should_stop()?;
+        df.rechunk_mut_par();
 
         let height = df.height();
 

crates/polars-mem-engine/src/planner/lp.rs

@@ -809,6 +809,8 @@ fn create_physical_plan_impl(
             input_left,
             input_right,
             key,
+            // In the in-memory engine, merge_sorted is always order-maintaining.
+            maintain_order: _,
         } => {
             let (input_left, input_right) = state.with_new_branch(|new_state| {
                 (

crates/polars-ops/src/frame/join/merge_sorted.rs

@@ -1,6 +1,10 @@
 use arrow::legacy::utils::CustomIterTools;
+#[cfg(feature = "dtype-categorical")]
+use polars_core::datatypes::CategoricalPhysical;
 use polars_core::prelude::*;
-use polars_core::{with_match_categorical_physical_type, with_match_physical_numeric_polars_type};
+#[cfg(feature = "dtype-categorical")]
+use polars_core::with_match_categorical_physical_type;
+use polars_core::with_match_physical_numeric_polars_type;
 
 pub fn _merge_sorted_dfs(
     left: &DataFrame,
@@ -185,6 +189,7 @@ where
 }
 
 fn series_to_merge_indicator(lhs: &Series, rhs: &Series) -> PolarsResult<Vec<bool>> {
+    #[cfg(feature = "dtype-categorical")]
     if lhs.dtype().is_categorical() || lhs.dtype().is_enum() {
         let cat_phys = lhs.dtype().cat_physical().unwrap();
         with_match_categorical_physical_type!(cat_phys, |$C| {

crates/polars-plan/dsl-schema-hashes.json

@@ -44,7 +44,7 @@
   "Dimension": "68880cdb10230df6c8c1632b073c80bd8ceb5c56a368c0cb438431ca9f3d3b31",
   "DistinctOptionsDSL": "41be5ec69ef9a614f2b36ac5deadfecdea5cca847ae1ada9d4bc626ff52a5b38",
   "DslFunction": "221f1a46a043c8ed54f57be981bf24509f04f5f91f0f08e0acc180d96f842ebf",
-  "DslPlan": "14caf5b73e69c4975ff3a57331891521ff5b78c96bbaf8d6cc9be57c82f3ea98",
+  "DslPlan": "037aeb1be892efd716c6934961e6df74dcd38815064b6d7efa72efe41e6e913d",
   "Duration": "44999d59023085cbb592ce94b30d34f9b983081fc72bd6435a49bdf0869c0074",
   "Duration2": "f251cb1bee2955a17c6defe1573bce21ddbe6cdf6eb9324a19cd37932ab29347",
   "DynListLiteralValue": "2266a553cb4a943f7097f24539eaa802453cf8742675996215235bd682dec0e8",

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

@@ -169,6 +169,7 @@ pub enum DslPlan {
         input_left: Arc<DslPlan>,
         input_right: Arc<DslPlan>,
         key: PlSmallStr,
+        maintain_order: bool,
     },
     IR {
         // Keep the original Dsl around as we need that for serialization.
@@ -211,7 +212,7 @@ impl Clone for DslPlan {
             #[cfg(feature = "pivot")]
             Self::Pivot { input, on, on_columns, index, values, agg, separator, maintain_order, column_naming }  => Self::Pivot { input: input.clone(), on: on.clone(), on_columns: on_columns.clone(), index: index.clone(), values: values.clone(), agg: agg.clone(), separator: separator.clone(), maintain_order: *maintain_order, column_naming: *column_naming },
             #[cfg(feature = "merge_sorted")]
-            Self::MergeSorted { input_left, input_right, key } => Self::MergeSorted { input_left: input_left.clone(), input_right: input_right.clone(), key: key.clone() },
+            Self::MergeSorted { input_left, input_right, key, maintain_order } => Self::MergeSorted { input_left: input_left.clone(), input_right: input_right.clone(), key: key.clone(), maintain_order: *maintain_order },
             Self::IR {node, dsl, version} => Self::IR {node: *node, dsl: dsl.clone(), version: *version},
         }
     }

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

@@ -146,6 +146,7 @@ pub(crate) enum SerializableDslPlanNode {
         input_left: DslPlanKey,
         input_right: DslPlanKey,
         key: PlSmallStr,
+        maintain_order: bool,
     },
     IR {
         dsl: DslPlanKey,
@@ -360,10 +361,12 @@ fn convert_dsl_plan_to_serializable_plan(
             input_left,
             input_right,
             key,
+            maintain_order,
         } => SP::MergeSorted {
             input_left: dsl_plan_key(input_left, arenas),
             input_right: dsl_plan_key(input_right, arenas),
             key: key.clone(),
+            maintain_order: *maintain_order,
         },
         DP::IR {
             dsl,
@@ -608,10 +611,12 @@ fn try_convert_serializable_plan_to_dsl_plan(
             input_left,
             input_right,
             key,
+            maintain_order,
         } => Ok(DP::MergeSorted {
             input_left: get_dsl_plan(*input_left, ser_dsl_plan, arenas)?,
             input_right: get_dsl_plan(*input_right, ser_dsl_plan, arenas)?,
             key: key.clone(),
+            maintain_order: *maintain_order,
         }),
         SP::IR {
             dsl: dsl_key,