Expand regex function input (#3613)

pola-rs · Jun 7, 2022 · 4fa8fad · 4fa8fad
1 parent 033b039
commit 4fa8fad
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 9 deletions.
diff --git a/polars/polars-lazy/src/dsl/functions.rs b/polars/polars-lazy/src/dsl/functions.rs
@@ -184,7 +184,8 @@ pub fn argsort_by<E: AsRef<[Expr]>>(by: E, reverse: &[bool]) -> Expr {
 #[cfg(feature = "concat_str")]
 #[cfg_attr(docsrs, doc(cfg(feature = "concat_str")))]
 /// Horizontally concat string columns in linear time
-pub fn concat_str(s: Vec<Expr>, sep: &str) -> Expr {
+pub fn concat_str<E: AsRef<[Expr]>>(s: E, sep: &str) -> Expr {
+    let s = s.as_ref().to_vec();
     let sep = sep.to_string();
     let function = NoEq::new(Arc::new(move |s: &mut [Series]| {
         polars_core::functions::concat_str(s, &sep).map(|ca| ca.into_series())

diff --git a/polars/polars-lazy/src/logical_plan/projection.rs b/polars/polars-lazy/src/logical_plan/projection.rs
@@ -245,7 +245,10 @@ 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::AnonymousFunction { input, options, .. } if options.input_wildcard_expansion => {
+            Expr::AnonymousFunction { input, options, .. }
+            | Expr::Function { input, options, .. }
+                if options.input_wildcard_expansion =>
+            {
                 if input
                     .iter()
                     .any(|e| matches!(e, Expr::Columns(_) | Expr::DtypeColumn(_)))
@@ -291,7 +294,10 @@ fn expand_function_list_inputs(mut expr: Expr, schema: &Schema) -> Expr {
 fn function_wildcard_expansion(mut expr: Expr, schema: &Schema, exclude: &[Arc<str>]) -> Expr {
     expr.mutate().apply(|e| {
         match e {
-            Expr::AnonymousFunction { input, options, .. } if options.input_wildcard_expansion => {
+            Expr::AnonymousFunction { input, options, .. }
+            | Expr::Function { input, options, .. }
+                if options.input_wildcard_expansion =>
+            {
                 let mut new_inputs = Vec::with_capacity(input.len());
 
                 input.iter_mut().for_each(|e| {
@@ -356,14 +362,16 @@ pub(crate) fn rewrite_projections(exprs: Vec<Expr>, schema: &Schema, keys: &[Exp
             replace_nth(&mut expr, schema);
         }
 
-        if has_wildcard(&expr) {
+        let function_input_expansion = has_expr(
+            &expr,
+            |e| matches!(e, Expr::AnonymousFunction { options,  .. } | Expr::Function {options, ..} if options.input_wildcard_expansion),
+        );
+
+        if has_wildcard(&expr) || function_input_expansion {
             // keep track of column excluded from the wildcard
             let exclude = prepare_excluded(&expr, schema, keys);
             // this path prepares the wildcard as input for the Function Expr
-            if has_expr(
-                &expr,
-                |e| matches!(e, Expr::AnonymousFunction { options,  .. } if options.input_wildcard_expansion),
-            ) {
+            if function_input_expansion {
                 expr = function_wildcard_expansion(expr, schema, &exclude);
                 result.push(expr);
                 continue;

diff --git a/polars/polars-lazy/src/tests/projection_queries.rs b/polars/polars-lazy/src/tests/projection_queries.rs
@@ -105,3 +105,21 @@ fn scan_join_same_file() -> Result<()> {
     );
     Ok(())
 }
+
+#[test]
+fn concat_str_regex_expansion() -> Result<()> {
+    let df = df![
+        "a"=> [1, 1, 1],
+        "b_a_1"=> ["a--", "", ""],
+        "b_a_2"=> ["", "b--", ""],
+        "b_a_3"=> ["", "", "c--"]
+    ]?
+    .lazy();
+    let out = df
+        .select([concat_str([col(r"^b_a_\d$")], ";").alias("concatenated")])
+        .collect()?;
+    let s = out.column("concatenated")?;
+    assert_eq!(s, &Series::new("concatenated", ["a--;;", ";b--;", ";;c--"]));
+
+    Ok(())
+}
diff --git a/py-polars/src/lib.rs b/py-polars/src/lib.rs
@@ -209,7 +209,7 @@ fn toggle_string_cache(toggle: bool) {
 
 #[pyfunction]
 fn concat_str(s: Vec<dsl::PyExpr>, sep: &str) -> dsl::PyExpr {
-    let s = s.into_iter().map(|e| e.inner).collect();
+    let s = s.into_iter().map(|e| e.inner).collect::<Vec<_>>();
     polars::lazy::dsl::concat_str(s, sep).into()
 }