fix: Check for MAP-GROUPS in cloud-eligible (#20662)

pola-rs · Jan 10, 2025 · b816b97 · b816b97
1 parent c4b704b
commit b816b97
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/crates/polars-plan/src/client/check.rs b/crates/polars-plan/src/client/check.rs
@@ -12,6 +12,9 @@ pub(super) fn assert_cloud_eligible(dsl: &DslPlan) -> PolarsResult<()> {
         match plan_node {
             #[cfg(feature = "python")]
             DslPlan::PythonScan { .. } => return ineligible_error("contains Python scan"),
+            DslPlan::GroupBy { apply, .. } if apply.is_some() => {
+                return ineligible_error("contains map groups")
+            },
             DslPlan::Scan {
                 sources, scan_type, ..
             } => {

diff --git a/py-polars/tests/unit/cloud/test_prepare_cloud_plan.py b/py-polars/tests/unit/cloud/test_prepare_cloud_plan.py
@@ -38,12 +38,6 @@ def test_prepare_cloud_plan(lf: pl.LazyFrame) -> None:
             pl.col("b").map_batches(lambda x: sum(x))
         ),
         pl.LazyFrame({"a": [1, 2], "b": [3, 4]}).map_batches(lambda x: x),
-        pl.LazyFrame({"a": [1, 2], "b": [3, 4]})
-        .group_by("a")
-        .map_groups(lambda x: x, schema={"b": pl.Int64}),
-        pl.LazyFrame({"a": [1, 2], "b": [3, 4]})
-        .group_by("a")
-        .agg(pl.col("b").map_batches(lambda x: sum(x))),
         pl.scan_parquet(CLOUD_SOURCE).filter(
             pl.col("a") < pl.lit(1).map_elements(lambda x: x + 1)
         ),