ARROW-6091: [Rust] [DataFusion] Implement physical execution plan for LIMIT

rust/arrow/src/array/builder.rs

@@ -1458,5 +1458,4 @@ mod tests {
         let mut builder = StructBuilder::new(fields, field_builders);
         assert!(builder.field_builder::<BinaryBuilder>(0).is_none());
     }
-
 }

rust/arrow/src/bitmap.rs

@@ -122,5 +122,4 @@ mod tests {
         assert_eq!(true, bitmap.is_set(6));
         assert_eq!(false, bitmap.is_set(7));
     }
-
 }

rust/datafusion/src/execution/aggregate.rs

@@ -1471,5 +1471,4 @@ mod tests {
             ds,
         )))))
     }
-
 }

rust/datafusion/src/execution/context.rs

@@ -682,5 +682,4 @@ mod tests {
 
         Ok(ctx)
     }
-
 }

rust/datafusion/src/execution/limit.rs

@@ -21,12 +21,11 @@ use std::cell::RefCell;
 use std::rc::Rc;
 use std::sync::Arc;
 
-use arrow::array::*;
-use arrow::compute::array_ops::limit;
 use arrow::datatypes::Schema;
 use arrow::record_batch::RecordBatch;
 
-use crate::error::{ExecutionError, Result};
+use crate::error::Result;
+use crate::execution::physical_plan::limit::truncate_batch;
 use crate::execution::relation::Relation;
 
 /// Implementation of a LIMIT relation
@@ -67,18 +66,9 @@ impl Relation for LimitRelation {
                     return Ok(None);
                 }
 
-                if batch.num_rows() >= capacity {
-                    let limited_columns: Result<Vec<ArrayRef>> = (0..batch.num_columns())
-                        .map(|i| match limit(batch.column(i), capacity) {
-                            Ok(result) => Ok(result),
-                            Err(error) => Err(ExecutionError::from(error)),
-                        })
-                        .collect();
-
-                    let limited_batch: RecordBatch =
-                        RecordBatch::try_new(self.schema.clone(), limited_columns?)?;
+                if batch.num_rows() > capacity {
+                    let limited_batch = truncate_batch(&batch, capacity)?;
                     self.num_consumed_rows += capacity;
-
                     Ok(Some(limited_batch))
                 } else {
                     self.num_consumed_rows += batch.num_rows();

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

@@ -720,5 +720,4 @@ mod tests {
 
         Ok(())
     }
-
 }

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

@@ -0,0 +1,208 @@
+// 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 LIMIT plan
+
+use crate::error::{ExecutionError, Result};
+use crate::execution::physical_plan::common::RecordBatchIterator;
+use crate::execution::physical_plan::ExecutionPlan;
+use crate::execution::physical_plan::{BatchIterator, Partition};
+use arrow::array::ArrayRef;
+use arrow::compute::array_ops::limit;
+use arrow::datatypes::Schema;
+use arrow::record_batch::RecordBatch;
+use std::sync::{Arc, Mutex};
+use std::thread;
+use std::thread::JoinHandle;
+
+/// Limit execution plan
+pub struct LimitExec {
+    /// Input schema
+    schema: Arc<Schema>,
+    /// Input partitions
+    partitions: Vec<Arc<dyn Partition>>,
+    /// Maximum number of rows to return
+    limit: usize,
+}
+
+impl LimitExec {
+    /// Create a new MergeExec
+    pub fn new(
+        schema: Arc<Schema>,
+        partitions: Vec<Arc<dyn Partition>>,
+        limit: usize,
+    ) -> Self {
+        LimitExec {
+            schema,
+            partitions,
+            limit,
+        }
+    }
+}
+
+impl ExecutionPlan for LimitExec {
+    fn schema(&self) -> Arc<Schema> {
+        self.schema.clone()
+    }
+
+    fn partitions(&self) -> Result<Vec<Arc<dyn Partition>>> {
+        Ok(vec![Arc::new(LimitPartition {
+            schema: self.schema.clone(),
+            partitions: self.partitions.clone(),
+            limit: self.limit,
+        })])
+    }
+}
+
+struct LimitPartition {
+    /// Input schema
+    schema: Arc<Schema>,
+    /// Input partitions
+    partitions: Vec<Arc<dyn Partition>>,
+    /// Maximum number of rows to return
+    limit: usize,
+}
+
+impl Partition for LimitPartition {
+    fn execute(&self) -> Result<Arc<Mutex<dyn BatchIterator>>> {
+        // collect up to "limit" rows on each partition
+        let threads: Vec<JoinHandle<Result<Vec<RecordBatch>>>> = self
+            .partitions
+            .iter()
+            .map(|p| {
+                let p = p.clone();
+                let limit = self.limit;
+                thread::spawn(move || {
+                    let it = p.execute()?;
+                    collect_with_limit(it, limit)
+                })
+            })
+            .collect();
+
+        // combine the results from each thread, up to the limit
+        let mut combined_results: Vec<Arc<RecordBatch>> = vec![];
+        let mut count = 0;
+        for thread in threads {
+            let join = thread.join().expect("Failed to join thread");
+            let result = join?;
+            for batch in result {
+                let capacity = self.limit - count;
+                if batch.num_rows() <= capacity {
+                    count += batch.num_rows();
+                    combined_results.push(Arc::new(batch.clone()))
+                } else {
+                    let batch = truncate_batch(&batch, capacity)?;
+                    count += batch.num_rows();
+                    combined_results.push(Arc::new(batch.clone()))
+                }
+                if count == self.limit {
+                    break;
+                }
+            }
+        }
+
+        Ok(Arc::new(Mutex::new(RecordBatchIterator::new(
+            self.schema.clone(),
+            combined_results,
+        ))))
+    }
+}
+
+/// Truncate a RecordBatch to maximum of n rows
+pub fn truncate_batch(batch: &RecordBatch, n: usize) -> Result<RecordBatch> {
+    let limited_columns: Result<Vec<ArrayRef>> = (0..batch.num_columns())
+        .map(|i| match limit(batch.column(i), n) {
+            Ok(result) => Ok(result),
+            Err(error) => Err(ExecutionError::from(error)),
+        })
+        .collect();
+
+    Ok(RecordBatch::try_new(
+        batch.schema().clone(),
+        limited_columns?,
+    )?)
+}
+
+/// Create a vector of record batches from an iterator
+fn collect_with_limit(
+    it: Arc<Mutex<dyn BatchIterator>>,
+    limit: usize,
+) -> Result<Vec<RecordBatch>> {
+    let mut count = 0;
+    let mut it = it.lock().unwrap();
+    let mut results: Vec<RecordBatch> = vec![];
+    loop {
+        match it.next() {
+            Ok(Some(batch)) => {
+                let capacity = limit - count;
+                if batch.num_rows() <= capacity {
+                    count += batch.num_rows();
+                    results.push(batch);
+                } else {
+                    let batch = truncate_batch(&batch, capacity)?;
+                    count += batch.num_rows();
+                    results.push(batch);
+                }
+                if count == limit {
+                    return Ok(results);
+                }
+            }
+            Ok(None) => {
+                // end of result set
+                return Ok(results);
+            }
+            Err(e) => return Err(e),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use super::*;
+    use crate::execution::physical_plan::common;
+    use crate::execution::physical_plan::csv::CsvExec;
+    use crate::test;
+
+    #[test]
+    fn limit() -> 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 limit = LimitExec::new(schema.clone(), input, 7);
+        let partitions = limit.partitions()?;
+
+        // the result should contain 4 batches (one per input partition)
+        let iter = partitions[0].execute()?;
+        let batches = common::collect(iter)?;
+
+        // there should be a total of 100 rows
+        let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum();
+        assert_eq!(row_count, 7);
+
+        Ok(())
+    }
+}

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

@@ -134,5 +134,4 @@ mod tests {
 
         Ok(())
     }
-
 }

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

@@ -86,6 +86,7 @@ pub mod csv;
 pub mod datasource;
 pub mod expressions;
 pub mod hash_aggregate;
+pub mod limit;
 pub mod merge;
 pub mod projection;
 pub mod selection;

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

@@ -171,5 +171,4 @@ mod tests {
 
         Ok(())
     }
-
 }

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

@@ -180,5 +180,4 @@ mod tests {
 
         Ok(())
     }
-
 }

rust/datafusion/src/execution/projection.rs

@@ -141,5 +141,4 @@ mod tests {
 
         assert_eq!("c1", batch.schema().field(0).name());
     }
-
 }

rust/datafusion/src/execution/table_impl.rs

@@ -315,5 +315,4 @@ mod tests {
             true,
         );
     }
-
 }

rust/datafusion/src/logicalplan.rs

@@ -641,5 +641,4 @@ mod tests {
             println!("plan: {:?}", plan1);
         });
     }
-
 }

rust/datafusion/src/sql/planner.rs

@@ -569,5 +569,4 @@ mod tests {
             }
         }
     }
-
 }

rust/parquet/src/compression.rs

@@ -338,5 +338,4 @@ mod tests {
     fn test_codec_zstd() {
         test_codec(CodecType::ZSTD);
     }
-
 }

rust/rust-toolchain

@@ -1 +1 @@
-nightly-2019-07-30
+nightly-2019-09-25