improve nested wildcard expansion (#2870)

pola-rs · Mar 10, 2022 · fe28511 · fe28511
1 parent 3fe3dc0
commit fe28511
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 22 deletions.
diff --git a/polars/polars-lazy/src/logical_plan/projection.rs b/polars/polars-lazy/src/logical_plan/projection.rs
@@ -267,28 +267,32 @@ pub(crate) fn rewrite_projections(exprs: Vec<Expr>, schema: &Schema, keys: &[Exp
                 |e| matches!(e, Expr::Function { options,  .. } if options.input_wildcard_expansion),
             ) {
                 expr.mutate().apply(|e| {
-                    if let Expr::Function { input, .. } = e {
-                        let mut new_inputs = Vec::with_capacity(input.len());
+                    match e {
+                        Expr::Function { input, options, .. }
+                            if options.input_wildcard_expansion =>
+                        {
+                            let mut new_inputs = Vec::with_capacity(input.len());
 
-                        input.iter_mut().for_each(|e| {
-                            if has_wildcard(e) {
-                                replace_wilcard(e, &mut new_inputs, &exclude, schema)
-                            } else {
-                                #[cfg(feature = "regex")]
-                                {
-                                    replace_regex(e, &mut new_inputs, schema)
-                                }
-                                #[cfg(not(feature = "regex"))]
-                                {
-                                    new_inputs.push(e.clone())
-                                }
-                            };
-                        });
+                            input.iter_mut().for_each(|e| {
+                                if has_wildcard(e) {
+                                    replace_wilcard(e, &mut new_inputs, &exclude, schema)
+                                } else {
+                                    #[cfg(feature = "regex")]
+                                    {
+                                        replace_regex(e, &mut new_inputs, schema)
+                                    }
+                                    #[cfg(not(feature = "regex"))]
+                                    {
+                                        new_inputs.push(e.clone())
+                                    }
+                                };
+                            });
 
-                        *input = new_inputs;
-                        false
-                    } else {
-                        true
+                            *input = new_inputs;
+                            // continue there can be more functions that require expansion
+                            true
+                        }
+                        _ => true,
                     }
                 });
                 result.push(expr);

diff --git a/polars/polars-lazy/src/physical_plan/expressions/column.rs b/polars/polars-lazy/src/physical_plan/expressions/column.rs
@@ -66,7 +66,12 @@ impl PhysicalExpr for ColumnExpr {
                     // now we do a linear search first as the lazy reported schema may still be incorrect
                     // in debug builds we panic so that it can be fixed when occurring
                     None => {
-                        debug_assert!(false);
+                        #[cfg(feature = "panic_on_schema")]
+                        {
+                            panic!("invalid schema")
+                        }
+                        // in release we fallback to linear search
+                        #[allow(unreachable_code)]
                         df.column(&self.0).map(|s| s.clone())
                     }
                 }

diff --git a/py-polars/polars/internals/lazy_functions.py b/py-polars/polars/internals/lazy_functions.py
@@ -1071,7 +1071,9 @@ def _date(
     return _datetime(year, month, day).cast(Date).alias("date")
 
 
-def concat_str(exprs: Sequence[Union["pli.Expr", str]], sep: str = "") -> "pli.Expr":
+def concat_str(
+    exprs: Union[Sequence[Union["pli.Expr", str]], "pli.Expr"], sep: str = ""
+) -> "pli.Expr":
     """
     Horizontally Concat Utf8 Series in linear time. Non utf8 columns are cast to utf8.
 

diff --git a/py-polars/tests/test_exprs.py b/py-polars/tests/test_exprs.py
@@ -111,3 +111,14 @@ def test_unique_stable() -> None:
     expected = pl.Series("a", [1, 2, 3])
 
     verify_series_and_expr_api(a, expected, "unique", True)
+
+
+def test_wildcard_expansion() -> None:
+    # one function requires wildcard expansion the other need
+    # this tests the nested behavior
+    # see: #2867
+
+    df = pl.DataFrame({"a": ["x", "Y", "z"], "b": ["S", "o", "S"]})
+    assert df.select(
+        pl.concat_str(pl.all()).str.to_lowercase()
+    ).to_series().to_list() == ["xs", "yo", "zs"]