ARROW-6563: [Rust] [DataFusion] MergeExec

rust/datafusion/src/execution/context.rs

@@ -22,8 +22,6 @@ use std::collections::HashMap;
 use std::rc::Rc;
 use std::string::String;
 use std::sync::Arc;
-use std::thread;
-use std::thread::JoinHandle;
 
 use arrow::datatypes::*;
 
@@ -37,8 +35,10 @@ use crate::execution::aggregate::AggregateRelation;
 use crate::execution::expression::*;
 use crate::execution::filter::FilterRelation;
 use crate::execution::limit::LimitRelation;
+use crate::execution::physical_plan::common;
 use crate::execution::physical_plan::datasource::DatasourceExec;
 use crate::execution::physical_plan::expressions::Column;
+use crate::execution::physical_plan::merge::MergeExec;
 use crate::execution::physical_plan::projection::ProjectionExec;
 use crate::execution::physical_plan::{ExecutionPlan, PhysicalExpr};
 use crate::execution::projection::ProjectRelation;
@@ -278,48 +278,28 @@ impl ExecutionContext {
 
     /// Execute a physical plan and collect the results in memory
     pub fn collect(&self, plan: &dyn ExecutionPlan) -> Result<Vec<RecordBatch>> {
-        let threads: Vec<JoinHandle<Result<Vec<RecordBatch>>>> = plan
-            .partitions()?
-            .iter()
-            .map(|p| {
-                let p = p.clone();
-                thread::spawn(move || {
-                    let it = p.execute()?;
-                    let mut it = it.lock().unwrap();
-                    let mut results: Vec<RecordBatch> = vec![];
-                    loop {
-                        match it.next() {
-                            Ok(Some(batch)) => {
-                                results.push(batch);
-                            }
-                            Ok(None) => {
-                                // end of result set
-                                return Ok(results);
-                            }
-                            Err(e) => return Err(e),
-                        }
-                    }
-                })
-            })
-            .collect();
-
-        // combine the results from each thread
-        let mut combined_results: Vec<RecordBatch> = vec![];
-        for thread in threads {
-            match thread.join() {
-                Ok(result) => {
-                    let result = result?;
-                    result
-                        .iter()
-                        .for_each(|batch| combined_results.push(batch.clone()));
-                }
-                Err(_) => {
-                    return Err(ExecutionError::General("Thread failed".to_string()))
+        let partitions = plan.partitions()?;
+
+        match partitions.len() {
+            0 => Ok(vec![]),
+            1 => {
+                let it = partitions[0].execute()?;
+                common::collect(it)
+            }
+            _ => {
+                // merge into a single partition
+                let plan = MergeExec::new(plan.schema().clone(), partitions);
+                let partitions = plan.partitions()?;
+                if partitions.len() == 1 {
+                    common::collect(partitions[0].execute()?)
+                } else {
+                    Err(ExecutionError::InternalError(format!(
+                        "MergeExec returned {} partitions",
+                        partitions.len()
+                    )))
                 }
             }
         }
-
-        Ok(combined_results)
     }
 
     /// Execute a logical plan and produce a Relation (a schema-aware iterator over a

rust/datafusion/src/execution/physical_plan/common.rs

@@ -0,0 +1,77 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Defines common code used in execution plans
+
+use std::sync::{Arc, Mutex};
+
+use crate::error::Result;
+use crate::execution::physical_plan::BatchIterator;
+
+use arrow::datatypes::Schema;
+use arrow::record_batch::RecordBatch;
+
+/// Iterator over a vector of record batches
+pub struct RecordBatchIterator {
+    schema: Arc<Schema>,
+    batches: Vec<Arc<RecordBatch>>,
+    index: usize,
+}
+
+impl RecordBatchIterator {
+    /// Create a new RecordBatchIterator
+    pub fn new(schema: Arc<Schema>, batches: Vec<Arc<RecordBatch>>) -> Self {
+        RecordBatchIterator {
+            schema,
+            index: 0,
+            batches,
+        }
+    }
+}
+
+impl BatchIterator for RecordBatchIterator {
+    fn schema(&self) -> Arc<Schema> {
+        self.schema.clone()
+    }
+
+    fn next(&mut self) -> Result<Option<RecordBatch>> {
+        if self.index < self.batches.len() {
+            self.index += 1;
+            Ok(Some(self.batches[self.index - 1].as_ref().clone()))
+        } else {
+            Ok(None)
+        }
+    }
+}
+
+/// Create a vector of record batches from an iterator
+pub fn collect(it: Arc<Mutex<dyn BatchIterator>>) -> Result<Vec<RecordBatch>> {
+    let mut it = it.lock().unwrap();
+    let mut results: Vec<RecordBatch> = vec![];
+    loop {
+        match it.next() {
+            Ok(Some(batch)) => {
+                results.push(batch);
+            }
+            Ok(None) => {
+                // end of result set
+                return Ok(results);
+            }
+            Err(e) => return Err(e),
+        }
+    }
+}

rust/datafusion/src/execution/physical_plan/merge.rs

@@ -0,0 +1,138 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Defines the merge plan for executing partitions in parallel and then merging the results
+//! into a single partition
+
+use crate::error::Result;
+use crate::execution::physical_plan::common::RecordBatchIterator;
+use crate::execution::physical_plan::{common, ExecutionPlan};
+use crate::execution::physical_plan::{BatchIterator, Partition};
+use arrow::datatypes::Schema;
+use arrow::record_batch::RecordBatch;
+use std::sync::{Arc, Mutex};
+use std::thread;
+use std::thread::JoinHandle;
+
+/// Merge execution plan executes partitions in parallel and combines them into a single
+/// partition. No guarantees are made about the order of the resulting partition.
+pub struct MergeExec {
+    /// Input schema
+    schema: Arc<Schema>,
+    /// Input partitions
+    partitions: Vec<Arc<dyn Partition>>,
+}
+
+impl MergeExec {
+    /// Create a new MergeExec
+    pub fn new(schema: Arc<Schema>, partitions: Vec<Arc<dyn Partition>>) -> Self {
+        MergeExec { schema, partitions }
+    }
+}
+
+impl ExecutionPlan for MergeExec {
+    fn schema(&self) -> Arc<Schema> {
+        self.schema.clone()
+    }
+
+    fn partitions(&self) -> Result<Vec<Arc<dyn Partition>>> {
+        Ok(vec![Arc::new(MergePartition {
+            schema: self.schema.clone(),
+            partitions: self.partitions.clone(),
+        })])
+    }
+}
+
+struct MergePartition {
+    /// Input schema
+    schema: Arc<Schema>,
+    /// Input partitions
+    partitions: Vec<Arc<dyn Partition>>,
+}
+
+impl Partition for MergePartition {
+    fn execute(&self) -> Result<Arc<Mutex<dyn BatchIterator>>> {
+        let threads: Vec<JoinHandle<Result<Vec<RecordBatch>>>> = self
+            .partitions
+            .iter()
+            .map(|p| {
+                let p = p.clone();
+                thread::spawn(move || {
+                    let it = p.execute()?;
+                    common::collect(it)
+                })
+            })
+            .collect();
+
+        // combine the results from each thread
+        let mut combined_results: Vec<Arc<RecordBatch>> = vec![];
+        for thread in threads {
+            let join = thread.join().expect("Failed to join thread");
+            let result = join?;
+            result
+                .iter()
+                .for_each(|batch| combined_results.push(Arc::new(batch.clone())));
+        }
+
+        Ok(Arc::new(Mutex::new(RecordBatchIterator::new(
+            self.schema.clone(),
+            combined_results,
+        ))))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use super::*;
+    use crate::execution::physical_plan::common;
+    use crate::execution::physical_plan::csv::CsvExec;
+    use crate::test;
+
+    #[test]
+    fn merge() -> Result<()> {
+        let schema = test::aggr_test_schema();
+
+        let num_partitions = 4;
+        let path =
+            test::create_partitioned_csv("aggregate_test_100.csv", num_partitions)?;
+
+        let csv = CsvExec::try_new(&path, schema.clone(), true, None, 1024)?;
+
+        // input should have 4 partitions
+        let input = csv.partitions()?;
+        assert_eq!(input.len(), num_partitions);
+
+        let merge = MergeExec::new(schema.clone(), input);
+
+        // output of MergeExec should have a single partition
+        let merged = merge.partitions()?;
+        assert_eq!(merged.len(), 1);
+
+        // the result should contain 4 batches (one per input partition)
+        let iter = merged[0].execute()?;
+        let batches = common::collect(iter)?;
+        assert_eq!(batches.len(), num_partitions);
+
+        // there should be a total of 100 rows
+        let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum();
+        assert_eq!(row_count, 100);
+
+        Ok(())
+    }
+
+}

rust/datafusion/src/execution/physical_plan/mod.rs

@@ -81,7 +81,9 @@ pub trait Accumulator {
     fn get_value(&self) -> Result<Option<ScalarValue>>;
 }
 
+pub mod common;
 pub mod csv;
 pub mod datasource;
 pub mod expressions;
+pub mod merge;
 pub mod projection;

rust/datafusion/src/execution/physical_plan/projection.rs

@@ -142,25 +142,10 @@ mod tests {
     use crate::execution::physical_plan::csv::CsvExec;
     use crate::execution::physical_plan::expressions::Column;
     use crate::test;
-    use arrow::datatypes::{DataType, Field, Schema};
 
     #[test]
     fn project_first_column() -> Result<()> {
-        let schema = Arc::new(Schema::new(vec![
-            Field::new("c1", DataType::Utf8, false),
-            Field::new("c2", DataType::UInt32, false),
-            Field::new("c3", DataType::Int8, false),
-            Field::new("c3", DataType::Int16, false),
-            Field::new("c4", DataType::Int32, false),
-            Field::new("c5", DataType::Int64, false),
-            Field::new("c6", DataType::UInt8, false),
-            Field::new("c7", DataType::UInt16, false),
-            Field::new("c8", DataType::UInt32, false),
-            Field::new("c9", DataType::UInt64, false),
-            Field::new("c10", DataType::Float32, false),
-            Field::new("c11", DataType::Float64, false),
-            Field::new("c12", DataType::Utf8, false),
-        ]));
+        let schema = test::aggr_test_schema();
 
         let partitions = 4;
         let path = test::create_partitioned_csv("aggregate_test_100.csv", partitions)?;