ARROW-6694: [Rust] [DataFusion] Integration tests now use physical query plan

rust/datafusion/src/execution/context.rs

@@ -41,6 +41,7 @@ use crate::execution::physical_plan::expressions::{
     Avg, BinaryExpr, CastExpr, Column, Count, Literal, Max, Min, Sum,
 };
 use crate::execution::physical_plan::hash_aggregate::HashAggregateExec;
+use crate::execution::physical_plan::limit::LimitExec;
 use crate::execution::physical_plan::merge::MergeExec;
 use crate::execution::physical_plan::projection::ProjectionExec;
 use crate::execution::physical_plan::selection::SelectionExec;
@@ -289,6 +290,43 @@ impl ExecutionContext {
                 let runtime_expr = self.create_physical_expr(expr, &input_schema)?;
                 Ok(Arc::new(SelectionExec::try_new(runtime_expr, input)?))
             }
+            LogicalPlan::Limit { input, expr, .. } => {
+                let input = self.create_physical_plan(input, batch_size)?;
+                let input_schema = input.as_ref().schema().clone();
+
+                match expr {
+                    &Expr::Literal(ref scalar_value) => {
+                        let limit: usize = match scalar_value {
+                            ScalarValue::Int8(limit) if *limit >= 0 => Ok(*limit as usize),
+                            ScalarValue::Int16(limit) if *limit >= 0 => {
+                                Ok(*limit as usize)
+                            }
+                            ScalarValue::Int32(limit) if *limit >= 0 => {
+                                Ok(*limit as usize)
+                            }
+                            ScalarValue::Int64(limit) if *limit >= 0 => {
+                                Ok(*limit as usize)
+                            }
+                            ScalarValue::UInt8(limit) => Ok(*limit as usize),
+                            ScalarValue::UInt16(limit) => Ok(*limit as usize),
+                            ScalarValue::UInt32(limit) => Ok(*limit as usize),
+                            ScalarValue::UInt64(limit) => Ok(*limit as usize),
+                            _ => Err(ExecutionError::ExecutionError(
+                                "Limit only supports non-negative integer literals"
+                                    .to_string(),
+                            )),
+                        }?;
+                        Ok(Arc::new(LimitExec::new(
+                            input_schema.clone(),
+                            input.partitions()?,
+                            limit,
+                        )))
+                    }
+                    _ => Err(ExecutionError::ExecutionError(
+                        "Limit only supports non-negative integer literals".to_string(),
+                    )),
+                }
+            }
             _ => Err(ExecutionError::General(
                 "Unsupported logical plan variant".to_string(),
             )),

rust/datafusion/tests/sql.rs

@@ -15,20 +15,18 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::cell::RefCell;
 use std::env;
-use std::rc::Rc;
 use std::sync::Arc;
 
 extern crate arrow;
 extern crate datafusion;
 
 use arrow::array::*;
 use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
 
 use datafusion::error::Result;
 use datafusion::execution::context::ExecutionContext;
-use datafusion::execution::relation::Relation;
 use datafusion::logicalplan::LogicalPlan;
 
 const DEFAULT_BATCH_SIZE: usize = 1024 * 1024;
@@ -101,9 +99,11 @@ fn parquet_single_nan_schema() {
     ctx.register_parquet("single_nan", &format!("{}/single_nan.parquet", testdata))
         .unwrap();
     let sql = "SELECT mycol FROM single_nan";
-    let relation = ctx.sql(&sql, 1024 * 1024).unwrap();
-    let mut results = relation.borrow_mut();
-    while let Some(batch) = results.next().unwrap() {
+    let plan = ctx.create_logical_plan(&sql).unwrap();
+    let plan = ctx.optimize(&plan).unwrap();
+    let plan = ctx.create_physical_plan(&plan, DEFAULT_BATCH_SIZE).unwrap();
+    let results = ctx.collect(plan.as_ref()).unwrap();
+    for batch in results {
         assert_eq!(1, batch.num_rows());
         assert_eq!(1, batch.num_columns());
     }
@@ -430,14 +430,14 @@ fn register_alltypes_parquet(ctx: &mut ExecutionContext) {
 fn execute(ctx: &mut ExecutionContext, sql: &str) -> String {
     let plan = ctx.create_logical_plan(&sql).unwrap();
     let plan = ctx.optimize(&plan).unwrap();
-    let results = ctx.execute(&plan, DEFAULT_BATCH_SIZE).unwrap();
+    let plan = ctx.create_physical_plan(&plan, DEFAULT_BATCH_SIZE).unwrap();
+    let results = ctx.collect(plan.as_ref()).unwrap();
     result_str(&results)
 }
 
-fn result_str(results: &Rc<RefCell<dyn Relation>>) -> String {
-    let mut relation = results.borrow_mut();
+fn result_str(results: &Vec<RecordBatch>) -> String {
     let mut str = String::new();
-    while let Some(batch) = relation.next().unwrap() {
+    for batch in results {
         for row_index in 0..batch.num_rows() {
             for column_index in 0..batch.num_columns() {
                 if column_index > 0 {