only expand function inputs if wildcard expansion allows it (#3039)

pola-rs · Apr 2, 2022 · 40815b0 · 40815b0
1 parent dec5b27
commit 40815b0
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/polars/polars-lazy/src/logical_plan/projection.rs b/polars/polars-lazy/src/logical_plan/projection.rs
@@ -240,7 +240,7 @@ fn prepare_excluded(expr: &Expr, schema: &Schema, keys: &[Expr]) -> Vec<Arc<str>
 fn expand_function_list_inputs(mut expr: Expr, schema: &Schema) -> Expr {
     expr.mutate().apply(|e| {
         match e {
-            Expr::Function { input, .. } => {
+            Expr::Function { input, options, .. } if options.input_wildcard_expansion => {
                 if input
                     .iter()
                     .any(|e| matches!(e, Expr::Columns(_) | Expr::DtypeColumn(_)))

diff --git a/polars/tests/it/lazy/expressions/apply.rs b/polars/tests/it/lazy/expressions/apply.rs
@@ -62,3 +62,23 @@ fn test_groups_update_binary_shift_log() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_expand_list() -> Result<()> {
+    let out = df![
+        "a" => [1, 2],
+        "b" => [2, 3],
+    ]?
+    .lazy()
+    .select([cols(["a", "b"]).cumsum(false)])
+    .collect()?;
+
+    let expected = df![
+        "a" => [1, 3],
+        "b" => [2, 5]
+    ]?;
+
+    assert!(out.frame_equal(&expected));
+
+    Ok(())
+}