feat(rust, python): allow regex and wildcard in groupby (#5425)

.github/deploy_manylinux.sh

@@ -8,7 +8,7 @@ ls -la
 rm py-polars/README.md
 cp README.md py-polars/README.md
 cd py-polars
-rustup override set nightly-2022-10-24
+rustup override set nightly-2022-11-03
 export RUSTFLAGS='-C target-feature=+fxsr,+sse,+sse2,+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt,+avx,+fma'
 
 # first the default release

.github/workflows/benchmark.yaml

@@ -50,7 +50,7 @@ jobs:
       - name: Set up Rust
         uses: dtolnay/rust-toolchain@master
         with:
-          toolchain: nightly-2022-10-24
+          toolchain: nightly-2022-11-03
 
       - name: Cache Rust
         uses: Swatinem/rust-cache@v2

.github/workflows/build-test-rust.yaml

@@ -34,7 +34,7 @@ jobs:
       - name: Set up Rust
         uses: dtolnay/rust-toolchain@master
         with:
-          toolchain: nightly-2022-10-24
+          toolchain: nightly-2022-11-03
 
       - name: Cache Rust
         uses: Swatinem/rust-cache@v2
@@ -58,7 +58,7 @@ jobs:
       - name: Set up Rust
         uses: dtolnay/rust-toolchain@master
         with:
-          toolchain: nightly-2022-10-24
+          toolchain: nightly-2022-11-03
           components: rustfmt, clippy, miri
 
       - name: Cache Rust
@@ -102,7 +102,7 @@ jobs:
       - name: Set up Rust
         uses: dtolnay/rust-toolchain@master
         with:
-          toolchain: nightly-2022-10-24
+          toolchain: nightly-2022-11-03
 
       - name: Cache Rust
         uses: Swatinem/rust-cache@v2

.github/workflows/create-py-release-mac-universal2.yaml

@@ -26,7 +26,7 @@ jobs:
       - name: Set up Rust
         uses: dtolnay/rust-toolchain@master
         with:
-          toolchain: nightly-2022-10-24
+          toolchain: nightly-2022-11-03
 
       - name: Set up Rust targets
         run: rustup target add aarch64-apple-darwin

.github/workflows/create-py-release-manylinux-lts-cpu.yaml

@@ -32,7 +32,7 @@ jobs:
           MATURIN_PASSWORD: ${{ secrets.PYPI_PASS }}
           RUSTFLAGS: -C target-feature=+fxsr,+sse,+sse2,+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt --cfg use_mimalloc
         with:
-          rust-toolchain: nightly-2022-10-24
+          rust-toolchain: nightly-2022-11-03
           maturin-version: '0.13.5'
           command: publish
           args: -m py-polars/Cargo.toml --skip-existing -o wheels -u ritchie46

.github/workflows/create-py-release-manylinux.yaml

@@ -29,7 +29,7 @@ jobs:
           MATURIN_PASSWORD: ${{ secrets.PYPI_PASS }}
           RUSTFLAGS: -C target-feature=+fxsr,+sse,+sse2,+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt,+avx,+fma
         with:
-          rust-toolchain: nightly-2022-10-24
+          rust-toolchain: nightly-2022-11-03
           maturin-version: '0.13.5'
           command: publish
           args: -m py-polars/Cargo.toml --skip-existing -o wheels -u ritchie46
@@ -53,7 +53,7 @@ jobs:
         env:
           MATURIN_PASSWORD: ${{ secrets.PYPI_PASS }}
         with:
-          rust-toolchain: nightly-2022-10-24
+          rust-toolchain: nightly-2022-11-03
           target: aarch64-unknown-linux-gnu
           maturin-version: '0.13.5'
           command: publish
@@ -84,7 +84,7 @@ jobs:
           MATURIN_PASSWORD: ${{ secrets.PYPI_PASS }}
           RUSTFLAGS: -C target-feature=+fxsr,+sse,+sse2,+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt,+avx,+fma
         with:
-          rust-toolchain: nightly-2022-10-24
+          rust-toolchain: nightly-2022-11-03
           maturin-version: '0.13.5'
           command: publish
           args: -m py-polars/Cargo.toml --skip-existing -o wheels -u ritchie46

.github/workflows/create-py-release-windows-macos.yaml

@@ -34,7 +34,7 @@ jobs:
           MATURIN_PASSWORD: ${{ secrets.PYPI_PASS }}
           RUSTFLAGS: -C target-feature=+fxsr,+sse,+sse2,+sse3,+sse4.1,+sse4.2
         with:
-          rust-toolchain: nightly-2022-10-24
+          rust-toolchain: nightly-2022-11-03
           maturin-version: '0.13.5'
           command: publish
           args: -m py-polars/Cargo.toml --no-sdist --skip-existing -o wheels -i python -u ritchie46

.github/workflows/docs-deploy.yaml

@@ -30,7 +30,7 @@ jobs:
       - name: Set up Rust
         uses: dtolnay/rust-toolchain@master
         with:
-          toolchain: nightly-2022-10-24
+          toolchain: nightly-2022-11-03
           components: rust-docs
 
       - name: Cache Rust

.github/workflows/lint-python.yaml

@@ -39,7 +39,7 @@ jobs:
       - name: Set up Rust
         uses: dtolnay/rust-toolchain@master
         with:
-          toolchain: nightly-2022-10-24
+          toolchain: nightly-2022-11-03
           components: rustfmt, clippy
       - name: Cache Rust
         uses: Swatinem/rust-cache@v2

.github/workflows/test-python.yaml

@@ -50,7 +50,7 @@ jobs:
       - name: Set up Rust
         uses: dtolnay/rust-toolchain@master
         with:
-          toolchain: nightly-2022-10-24
+          toolchain: nightly-2022-11-03
 
       - name: Cache Rust
         uses: Swatinem/rust-cache@v2
@@ -112,7 +112,7 @@ jobs:
       - name: Set up Rust
         uses: dtolnay/rust-toolchain@master
         with:
-          toolchain: nightly-2022-10-24
+          toolchain: nightly-2022-11-03
 
       - name: Cache Rust
         uses: Swatinem/rust-cache@v2

polars/polars-core/src/frame/cross_join.rs

@@ -48,7 +48,10 @@ impl DataFrame {
     ) -> PolarsResult<(DataFrame, DataFrame)> {
         let n_rows_left = self.height() as IdxSize;
         let n_rows_right = other.height() as IdxSize;
-        let total_rows = n_rows_right * n_rows_left;
+        let Some(total_rows) = n_rows_left.checked_mul(n_rows_right) else {
+            return Err(PolarsError::ComputeError("Cross joins would produce more rows than fits into 2^32.\n\
+            Consider comping with polars-big-idx feature, or set 'allow_streaming'.".into()))
+        };
 
         // the left side has the Nth row combined with every row from right.
         // So let's say we have the following no. of rows

polars/polars-lazy/polars-plan/src/logical_plan/builder.rs

@@ -365,7 +365,7 @@ impl LogicalPlanBuilder {
 
     pub fn groupby<E: AsRef<[Expr]>>(
         self,
-        keys: Arc<Vec<Expr>>,
+        keys: Vec<Expr>,
         aggs: E,
         apply: Option<Arc<dyn DataFrameUdf>>,
         maintain_order: bool,
@@ -374,6 +374,7 @@ impl LogicalPlanBuilder {
     ) -> Self {
         let current_schema = try_delayed!(self.0.schema(), &self.0, into);
         let current_schema = current_schema.as_ref();
+        let keys = rewrite_projections(keys, current_schema, &[]);
         let aggs = rewrite_projections(aggs.as_ref().to_vec(), current_schema, keys.as_ref());
 
         let mut schema = try_delayed!(
@@ -422,7 +423,7 @@ impl LogicalPlanBuilder {
 
         LogicalPlan::Aggregate {
             input: Box::new(self.0),
-            keys,
+            keys: Arc::new(keys),
             aggs,
             schema: Arc::new(schema),
             apply,

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

@@ -1249,7 +1249,7 @@ impl LazyGroupBy {
         #[cfg(feature = "dynamic_groupby")]
         let lp = LogicalPlanBuilder::from(self.logical_plan)
             .groupby(
-                Arc::new(self.keys),
+                self.keys,
                 aggs,
                 None,
                 self.maintain_order,
@@ -1260,7 +1260,7 @@ impl LazyGroupBy {
 
         #[cfg(not(feature = "dynamic_groupby"))]
         let lp = LogicalPlanBuilder::from(self.logical_plan)
-            .groupby(Arc::new(self.keys), aggs, None, self.maintain_order)
+            .groupby(self.keys, aggs, None, self.maintain_order)
             .build();
         LazyFrame::from_logical_plan(lp, self.opt_state)
     }

py-polars/src/dataframe.rs

@@ -948,7 +948,7 @@ impl PyDataFrame {
     }
 
     pub fn select(&self, selection: Vec<&str>) -> PyResult<Self> {
-        let df = self.df.select(&selection).map_err(PyPolarsErr::from)?;
+        let df = self.df.select(selection).map_err(PyPolarsErr::from)?;
         Ok(PyDataFrame::new(df))
     }
 

py-polars/src/file.rs

@@ -238,7 +238,7 @@ pub fn get_mmap_bytes_reader<'a>(py_f: &'a PyAny) -> PyResult<Box<dyn MmapBytesR
         let s = pstring.to_string();
         let p = std::path::Path::new(&s);
         let p = resolve_homedir(p);
-        let f = match File::open(&p) {
+        let f = match File::open(p) {
             Ok(file) => file,
             Err(_e) => {
                 return Err(PyErr::new::<PyFileNotFoundError, _>(format!(

py-polars/src/lazy/dsl.rs

@@ -1099,7 +1099,7 @@ impl PyExpr {
             .into()
     }
     pub fn exclude(&self, columns: Vec<String>) -> PyExpr {
-        self.inner.clone().exclude(&columns).into()
+        self.inner.clone().exclude(columns).into()
     }
     pub fn exclude_dtype(&self, dtypes: Vec<Wrap<DataType>>) -> PyExpr {
         // Safety:

py-polars/tests/unit/test_groupby.py

@@ -200,3 +200,15 @@ def test_take_in_groupby() -> None:
     assert df.groupby("group").agg(
         pl.col("values").take(1) - pl.col("values").take(2)
     ).sort("group").to_dict(False) == {"group": [1, 2], "values": [197, 494]}
+
+
+def test_groupby_wildcard() -> None:
+    df = pl.DataFrame(
+        {
+            "a": [1, 2],
+            "b": [1, 2],
+        }
+    )
+    assert df.groupby([pl.col("*")], maintain_order=True).agg(
+        [pl.col("a").first().suffix("_agg")]
+    ).to_dict(False) == {"a": [1, 2], "b": [1, 2], "a_agg": [1, 2]}

rust-toolchain.toml

@@ -1,2 +1,2 @@
 [toolchain]
-channel = "nightly-2022-10-24"
+channel = "nightly-2022-11-03"