feat: Improve scalar strict message (#18904)

pola-rs · Sep 24, 2024 · 77137ec · 77137ec
1 parent 0f4360b
commit 77137ec
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/crates/polars-mem-engine/src/executors/stack.rs b/crates/polars-mem-engine/src/executors/stack.rs
@@ -79,11 +79,19 @@ impl StackExec {
                 for (i, c) in res.iter().enumerate() {
                     let len = c.len();
                     if verify_scalar && len != df_height && len == 1 && df_width > 0 {
-                        polars_ensure!(self.exprs[i].is_scalar(),
-                            InvalidOperation: "Series {}, length {} doesn't match the DataFrame height of {}\n\n\
-                            If you want this Series to be broadcasted, ensure it is a scalar (for instance by adding '.first()').",
-                            c.name(), len, df_height
-                        );
+                        #[allow(clippy::collapsible_if)]
+                        if !self.exprs[i].is_scalar()
+                            && std::env::var("POLARS_ALLOW_NON_SCALAR_EXP").as_deref() != Ok("1")
+                        {
+                            let identifier = match self.exprs[i].as_expression() {
+                                Some(e) => format!("expression: {}", e),
+                                None => "this Series".to_string(),
+                            };
+                            polars_bail!(InvalidOperation: "Series {}, length {} doesn't match the DataFrame height of {}\n\n\
+                                If you want {} to be broadcasted, ensure it is a scalar (for instance by adding '.first()').",
+                                c.name(), len, df_height, identifier
+                            );
+                        }
                     }
                 }
                 df._add_series(res, schema)?;