From ed753fd6e4a1f0647225ccc3eb18ee77ae44b9ab Mon Sep 17 00:00:00 2001
From: Andy Grove <andygrove73@gmail.com>
Date: Sat, 5 Oct 2019 17:11:32 +0200
Subject: [PATCH] ARROW-6692: [Rust] [DataFusion] Update examples to use
 physical query plan
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Closes #5581 from andygrove/ARROW-6692 and squashes the following commits:

52e8da544 <Andy Grove> Update examples to use physical query plan

Authored-by: Andy Grove <andygrove73@gmail.com>
Signed-off-by: Krisztián Szűcs <szucs.krisztian@gmail.com>
---
 rust/datafusion/examples/csv_sql.rs     | 21 +++++++++++++--------
 rust/datafusion/examples/parquet_sql.rs | 23 ++++++++++++++---------
 2 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/rust/datafusion/examples/csv_sql.rs b/rust/datafusion/examples/csv_sql.rs
index e167f038701ca..9681af05db9d3 100644
--- a/rust/datafusion/examples/csv_sql.rs
+++ b/rust/datafusion/examples/csv_sql.rs
@@ -23,11 +23,12 @@ extern crate datafusion;
 use arrow::array::{BinaryArray, Float64Array};
 use arrow::datatypes::{DataType, Field, Schema};
 
+use datafusion::error::Result;
 use datafusion::execution::context::ExecutionContext;
 
 /// This example demonstrates executing a simple query against an Arrow data source (CSV) and
 /// fetching results
-fn main() {
+fn main() -> Result<()> {
     // create local execution context
     let mut ctx = ExecutionContext::new();
 
@@ -59,16 +60,18 @@ fn main() {
         true,
     );
 
-    // simple projection and selection
     let sql = "SELECT c1, MIN(c12), MAX(c12) FROM aggregate_test_100 WHERE c11 > 0.1 AND c11 < 0.9 GROUP BY c1";
 
-    // execute the query
-    let relation = ctx.sql(&sql, 1024 * 1024).unwrap();
+    // create the query plan
+    let plan = ctx.create_logical_plan(&sql)?;
+    let plan = ctx.optimize(&plan)?;
+    let plan = ctx.create_physical_plan(&plan, 1024 * 1024)?;
 
-    // display the relation
-    let mut results = relation.borrow_mut();
+    // execute the query
+    let results = ctx.collect(plan.as_ref())?;
 
-    while let Some(batch) = results.next().unwrap() {
+    // iterate over the results
+    results.iter().for_each(|batch| {
         println!(
             "RecordBatch has {} rows and {} columns",
             batch.num_rows(),
@@ -98,5 +101,7 @@ fn main() {
 
             println!("{}, Min: {}, Max: {}", c1_value, min.value(i), max.value(i),);
         }
-    }
+    });
+
+    Ok(())
 }
diff --git a/rust/datafusion/examples/parquet_sql.rs b/rust/datafusion/examples/parquet_sql.rs
index 81654441224b1..ec3924d94fec4 100644
--- a/rust/datafusion/examples/parquet_sql.rs
+++ b/rust/datafusion/examples/parquet_sql.rs
@@ -20,11 +20,12 @@ extern crate datafusion;
 
 use arrow::array::{BinaryArray, Float64Array, Int32Array};
 
+use datafusion::error::Result;
 use datafusion::execution::context::ExecutionContext;
 
 /// This example demonstrates executing a simple query against an Arrow data source (Parquet) and
 /// fetching results
-fn main() {
+fn main() -> Result<()> {
     // create local execution context
     let mut ctx = ExecutionContext::new();
 
@@ -35,19 +36,21 @@ fn main() {
     ctx.register_parquet(
         "alltypes_plain",
         &format!("{}/alltypes_plain.parquet", testdata),
-    )
-    .unwrap();
+    )?;
 
     // simple selection
     let sql = "SELECT int_col, double_col, date_string_col FROM alltypes_plain WHERE id > 1 AND tinyint_col < double_col";
 
-    // execute the query
-    let relation = ctx.sql(&sql, 1024 * 1024).unwrap();
+    // create the query plan
+    let plan = ctx.create_logical_plan(&sql)?;
+    let plan = ctx.optimize(&plan)?;
+    let plan = ctx.create_physical_plan(&plan, 1024 * 1024)?;
 
-    // display the relation
-    let mut results = relation.borrow_mut();
+    // execute the query
+    let results = ctx.collect(plan.as_ref())?;
 
-    while let Some(batch) = results.next().unwrap() {
+    // iterate over the results
+    results.iter().for_each(|batch| {
         println!(
             "RecordBatch has {} rows and {} columns",
             batch.num_rows(),
@@ -82,5 +85,7 @@ fn main() {
                 double.value(i)
             );
         }
-    }
+    });
+
+    Ok(())
 }