make lazy union parallel (#2221)

pola-rs · Dec 29, 2021 · 47f544d · 47f544d
1 parent 57074e5
commit 47f544d
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/polars/polars-lazy/src/physical_plan/executors/union.rs b/polars/polars-lazy/src/physical_plan/executors/union.rs
@@ -2,18 +2,24 @@ use crate::physical_plan::state::ExecutionState;
 use crate::prelude::*;
 use polars_core::prelude::*;
 use polars_core::utils::concat_df;
+use polars_core::POOL;
+use rayon::prelude::*;
 
 pub(crate) struct UnionExec {
     pub(crate) inputs: Vec<Box<dyn Executor>>,
 }
 
 impl Executor for UnionExec {
     fn execute(&mut self, state: &ExecutionState) -> Result<DataFrame> {
-        let dfs = self
-            .inputs
-            .iter_mut()
-            .map(|input| input.execute(state))
-            .collect::<Result<Vec<_>>>()?;
+        let inputs = std::mem::take(&mut self.inputs);
+
+        let dfs = POOL.install(|| {
+            inputs
+                .into_par_iter()
+                .map(|mut input| input.execute(state))
+                .collect::<Result<Vec<_>>>()
+        })?;
+
         concat_df(&dfs)
     }
 }