ARROW-7787: [Rust] Added .collect to Table API

This commit adds the method ".collect" to Table, thus
reducing the code necessary to execute a query constructed
from the Table API.

This also adds an example to the examples on how to use the Table API on a table declared in-Memory.

Closes #6375 from jorgecarleitao/table_collect and squashes the following commits:

8a6efaa <Jorge C. Leitao> ARROW-7787:  Added .collect to Table API

Authored-by: Jorge C. Leitao <jorgecarleitao@gmail.com>
Signed-off-by: Andy Grove <andygrove73@gmail.com>

rust/datafusion/examples/memory_table_api.rs

@@ -0,0 +1,97 @@
+// 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.
+
+use std::boxed::Box;
+use std::sync::Arc;
+
+extern crate arrow;
+extern crate datafusion;
+
+use arrow::array::{Int32Array, StringArray};
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+
+use datafusion::datasource::MemTable;
+use datafusion::execution::context::ExecutionContext;
+use datafusion::logicalplan::{Expr, ScalarValue};
+
+/// This example demonstrates basic uses of the Table API on an in-memory table
+fn main() {
+    // define a schema.
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("a", DataType::Utf8, false),
+        Field::new("b", DataType::Int32, false),
+    ]));
+
+    // define data.
+    let batch = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(StringArray::from(vec!["a", "b", "c", "d"])),
+            Arc::new(Int32Array::from(vec![1, 10, 10, 100])),
+        ],
+    )
+    .unwrap();
+
+    // declare a new context. In spark API, this corresponds to a new spark SQLsession
+    let mut ctx = ExecutionContext::new();
+
+    // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
+    let provider = MemTable::new(schema, vec![batch]).unwrap();
+    ctx.register_table("t", Box::new(provider));
+    let t = ctx.table("t").unwrap();
+
+    // construct an expression corresponding to "SELECT a, b FROM t WHERE b = 10" in SQL
+    let filter = t
+        .col("b")
+        .unwrap()
+        .eq(&Expr::Literal(ScalarValue::Int32(10)));
+
+    let t = t
+        .select_columns(vec!["a", "b"])
+        .unwrap()
+        .filter(filter)
+        .unwrap();
+
+    // execute
+    let results = t.collect(&mut ctx, 10).unwrap();
+
+    // print results
+    results.iter().for_each(|batch| {
+        println!(
+            "RecordBatch has {} rows and {} columns",
+            batch.num_rows(),
+            batch.num_columns()
+        );
+
+        let c1 = batch
+            .column(0)
+            .as_any()
+            .downcast_ref::<StringArray>()
+            .expect("String type");
+
+        let c2 = batch
+            .column(1)
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .expect("Int type");
+
+        for i in 0..batch.num_rows() {
+            println!("{}, {}", c1.value(i), c2.value(i),);
+        }
+    });
+}

rust/datafusion/src/execution/context.rs

@@ -71,7 +71,17 @@ impl ExecutionContext {
     pub fn sql(&mut self, sql: &str, batch_size: usize) -> Result<Vec<RecordBatch>> {
         let plan = self.create_logical_plan(sql)?;
 
-        match plan.as_ref() {
+        return self.collect_plan(plan.as_ref(), batch_size);
+    }
+
+    /// Executes a logical plan and produce a Relation (a schema-aware iterator over a series
+    /// of RecordBatch instances)
+    pub fn collect_plan(
+        &mut self,
+        plan: &LogicalPlan,
+        batch_size: usize,
+    ) -> Result<Vec<RecordBatch>> {
+        match plan {
             LogicalPlan::CreateExternalTable {
                 ref schema,
                 ref name,

rust/datafusion/src/execution/table_impl.rs

@@ -20,7 +20,9 @@
 use std::sync::Arc;
 
 use crate::arrow::datatypes::{DataType, Field, Schema};
+use crate::arrow::record_batch::RecordBatch;
 use crate::error::{ExecutionError, Result};
+use crate::execution::context::ExecutionContext;
 use crate::logicalplan::Expr::Literal;
 use crate::logicalplan::ScalarValue;
 use crate::logicalplan::{Expr, LogicalPlan};
@@ -154,6 +156,14 @@ impl Table for TableImpl {
     fn to_logical_plan(&self) -> Arc<LogicalPlan> {
         self.plan.clone()
     }
+
+    fn collect(
+        &self,
+        ctx: &mut ExecutionContext,
+        batch_size: usize,
+    ) -> Result<Vec<RecordBatch>> {
+        ctx.collect_plan(&self.plan.clone(), batch_size)
+    }
 }
 
 impl TableImpl {

rust/datafusion/src/table.rs

@@ -18,7 +18,9 @@
 //! Table API for building a logical query plan. This is similar to the Table API in Ibis
 //! and the DataFrame API in Apache Spark
 
+use crate::arrow::record_batch::RecordBatch;
 use crate::error::Result;
+use crate::execution::context::ExecutionContext;
 use crate::logicalplan::{Expr, LogicalPlan};
 use std::sync::Arc;
 
@@ -66,4 +68,11 @@ pub trait Table {
 
     /// Return the index of a column within this table's schema
     fn column_index(&self, name: &str) -> Result<usize>;
+
+    /// Collects the result as a vector of RecordBatch.
+    fn collect(
+        &self,
+        ctx: &mut ExecutionContext,
+        batch_size: usize,
+    ) -> Result<Vec<RecordBatch>>;
 }