[arrow2] Merge arrow2 and datafusion latest

Cargo.toml

@@ -35,6 +35,6 @@ lto = true
 codegen-units = 1
 
 [patch.crates-io]
-arrow2 = { git = "https://github.com/jorgecarleitao/arrow2.git", rev = "ef7937dfe56033c2cc491482c67587b52cd91554" }
+arrow2 = { git = "https://github.com/jorgecarleitao/arrow2.git", branch = "main" }
 #arrow2 = { git = "https://github.com/blaze-init/arrow2.git", branch = "shuffle_ipc" }
 #parquet2 = { git = "https://github.com/blaze-init/parquet2.git", branch = "meta_new" }

ballista/rust/client/src/columnar_batch.rs

@@ -23,8 +23,9 @@ use datafusion::arrow::{
     array::ArrayRef,
     compute::aggregate::estimated_bytes_size,
     datatypes::{DataType, Schema},
-    record_batch::RecordBatch,
 };
+use datafusion::field_util::{FieldExt, SchemaExt};
+use datafusion::record_batch::RecordBatch;
 use datafusion::scalar::ScalarValue;
 
 pub type MaybeColumnarBatch = Result<Option<ColumnarBatch>>;
@@ -44,7 +45,7 @@ impl ColumnarBatch {
             .enumerate()
             .map(|(i, array)| {
                 (
-                    batch.schema().field(i).name().clone(),
+                    batch.schema().field(i).name().to_string(),
                     ColumnarValue::Columnar(array.clone()),
                 )
             })
@@ -61,7 +62,7 @@ impl ColumnarBatch {
             .fields()
             .iter()
             .enumerate()
-            .map(|(i, f)| (f.name().clone(), values[i].clone()))
+            .map(|(i, f)| (f.name().to_string(), values[i].clone()))
             .collect();
 
         Self {

ballista/rust/core/Cargo.toml

@@ -43,8 +43,8 @@ tonic = "0.6"
 uuid = { version = "0.8", features = ["v4"] }
 chrono = { version = "0.4", default-features = false }
 
-arrow-format = { version = "0.3", features = ["flight-data", "flight-service"] }
-arrow = { package = "arrow2", version="0.8", features = ["io_ipc", "io_flight"] }
+arrow-format = { version = "0.4", features = ["flight-data", "flight-service"] }
+arrow = { package = "arrow2", version="0.9", features = ["io_ipc", "io_flight"] }
 
 datafusion = { path = "../../../datafusion", version = "6.0.0" }
 

ballista/rust/core/src/client.rs

@@ -39,10 +39,11 @@ use datafusion::arrow::{
     array::{StructArray, Utf8Array},
     datatypes::{Schema, SchemaRef},
     error::{ArrowError, Result as ArrowResult},
-    record_batch::RecordBatch,
 };
+use datafusion::field_util::SchemaExt;
 use datafusion::physical_plan::common::collect;
 use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream};
+use datafusion::record_batch::RecordBatch;
 use datafusion::{logical_plan::LogicalPlan, physical_plan::RecordBatchStream};
 use futures::{Stream, StreamExt};
 use log::debug;
@@ -174,11 +175,12 @@ impl Stream for FlightDataStream {
 
                         arrow::io::flight::deserialize_batch(
                             &flight_data_chunk,
-                            self.schema.clone(),
+                            self.schema.fields(),
                             &self.ipc_schema,
                             &hm,
                         )
-                    });
+                    })
+                    .map(|c| RecordBatch::new_with_chunk(&self.schema, c));
                 Some(converted_chunk)
             }
             None => None,

ballista/rust/core/src/execution_plans/shuffle_reader.rs

@@ -27,13 +27,13 @@ use crate::utils::WrappedStream;
 use async_trait::async_trait;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::arrow::error::Result as ArrowResult;
-use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::physical_plan::metrics::{
     ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
 };
 use datafusion::physical_plan::{
     DisplayFormatType, ExecutionPlan, Metric, Partitioning, Statistics,
 };
+use datafusion::record_batch::RecordBatch;
 use datafusion::{
     error::{DataFusionError, Result},
     physical_plan::RecordBatchStream,

ballista/rust/core/src/execution_plans/shuffle_writer.rs

@@ -33,6 +33,7 @@ use crate::utils;
 
 use crate::serde::protobuf::ShuffleWritePartition;
 use crate::serde::scheduler::{PartitionLocation, PartitionStats};
+use arrow::chunk::Chunk;
 use arrow::io::ipc::write::WriteOptions;
 use async_trait::async_trait;
 use datafusion::arrow::array::*;
@@ -41,8 +42,8 @@ use datafusion::arrow::compute::take;
 use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion::arrow::io::ipc::read::FileReader;
 use datafusion::arrow::io::ipc::write::FileWriter;
-use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::error::{DataFusionError, Result};
+use datafusion::field_util::SchemaExt;
 use datafusion::physical_plan::hash_utils::create_hashes;
 use datafusion::physical_plan::metrics::{
     self, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
@@ -52,6 +53,7 @@ use datafusion::physical_plan::Partitioning::RoundRobinBatch;
 use datafusion::physical_plan::{
     DisplayFormatType, ExecutionPlan, Metric, Partitioning, RecordBatchStream, Statistics,
 };
+use datafusion::record_batch::RecordBatch;
 use futures::StreamExt;
 use hashbrown::HashMap;
 use log::{debug, info};
@@ -468,7 +470,8 @@ impl ShuffleWriter {
     }
 
     fn write(&mut self, batch: &RecordBatch) -> Result<()> {
-        self.writer.write(batch, None)?;
+        let chunk = Chunk::new(batch.columns().to_vec());
+        self.writer.write(&chunk, None)?;
         self.num_batches += 1;
         self.num_rows += batch.num_rows() as u64;
         let num_bytes: usize = batch

ballista/rust/core/src/memory_stream.rs

@@ -20,8 +20,9 @@
 
 use std::task::{Context, Poll};
 
-use datafusion::arrow::{datatypes::SchemaRef, error::Result, record_batch::RecordBatch};
+use datafusion::arrow::{datatypes::SchemaRef, error::Result};
 use datafusion::physical_plan::RecordBatchStream;
+use datafusion::record_batch::RecordBatch;
 use futures::Stream;
 
 /// Iterator over batches

ballista/rust/core/src/serde/logical_plan/from_proto.rs

@@ -1169,6 +1169,7 @@ impl TryInto<Field> for &protobuf::Field {
 }
 
 use crate::serde::protobuf::ColumnStats;
+use datafusion::field_util::SchemaExt;
 use datafusion::physical_plan::{aggregates, windows};
 use datafusion::prelude::{
     array, date_part, date_trunc, length, lower, ltrim, md5, rtrim, sha224, sha256,

ballista/rust/core/src/serde/logical_plan/mod.rs

@@ -25,6 +25,7 @@ mod roundtrip_tests {
     use crate::error::BallistaError;
     use arrow::datatypes::UnionMode;
     use core::panic;
+    use datafusion::field_util::SchemaExt;
     use datafusion::logical_plan::Repartition;
     use datafusion::{
         arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit},

ballista/rust/core/src/serde/logical_plan/to_proto.rs

@@ -32,6 +32,7 @@ use datafusion::datasource::TableProvider;
 
 use datafusion::datasource::file_format::parquet::ParquetFormat;
 use datafusion::datasource::listing::ListingTable;
+use datafusion::field_util::{FieldExt, SchemaExt};
 use datafusion::logical_plan::plan::{
     Aggregate, EmptyRelation, Filter, Join, Projection, Sort, Window,
 };

ballista/rust/core/src/serde/physical_plan/mod.rs

@@ -22,6 +22,7 @@ pub mod to_proto;
 mod roundtrip_tests {
     use std::{convert::TryInto, sync::Arc};
 
+    use datafusion::field_util::SchemaExt;
     use datafusion::{
         arrow::{
             compute::sort::SortOptions,

ballista/rust/core/src/serde/physical_plan/to_proto.rs

@@ -26,6 +26,7 @@ use std::{
     sync::Arc,
 };
 
+use datafusion::field_util::FieldExt;
 use datafusion::physical_plan::hash_join::{HashJoinExec, PartitionMode};
 use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion::physical_plan::projection::ProjectionExec;
@@ -197,7 +198,7 @@ impl TryInto<protobuf::PhysicalPlanNode> for Arc<dyn ExecutionPlan> {
                 .aggr_expr()
                 .iter()
                 .map(|expr| match expr.field() {
-                    Ok(field) => Ok(field.name().clone()),
+                    Ok(field) => Ok(field.name().to_string()),
                     Err(e) => Err(BallistaError::DataFusionError(e)),
                 })
                 .collect::<Result<_, Self::Error>>()?;

ballista/rust/core/src/utils.rs

@@ -30,6 +30,7 @@ use crate::memory_stream::MemoryStream;
 use crate::serde::scheduler::PartitionStats;
 
 use crate::config::BallistaConfig;
+use arrow::chunk::Chunk;
 use async_trait::async_trait;
 use datafusion::arrow::datatypes::Schema;
 use datafusion::arrow::datatypes::SchemaRef;
@@ -41,12 +42,12 @@ use datafusion::arrow::{
     datatypes::{DataType, Field},
     io::ipc::read::FileReader,
     io::ipc::write::FileWriter,
-    record_batch::RecordBatch,
 };
 use datafusion::error::DataFusionError;
 use datafusion::execution::context::{
     ExecutionConfig, ExecutionContext, ExecutionContextState, QueryPlanner,
 };
+use datafusion::field_util::SchemaExt;
 use datafusion::logical_plan::{LogicalPlan, Operator};
 use datafusion::physical_optimizer::coalesce_batches::CoalesceBatches;
 use datafusion::physical_optimizer::merge_exec::AddCoalescePartitionsExec;
@@ -64,6 +65,7 @@ use datafusion::physical_plan::sort::SortExec;
 use datafusion::physical_plan::{
     metrics, AggregateExpr, ExecutionPlan, Metric, PhysicalExpr, RecordBatchStream,
 };
+use datafusion::record_batch::RecordBatch;
 use futures::{future, Stream, StreamExt};
 use std::time::Instant;
 
@@ -104,7 +106,8 @@ pub async fn write_stream_to_disk(
         num_bytes += batch_size_bytes;
 
         let timer = disk_write_metric.timer();
-        writer.write(&batch, None)?;
+        let chunk = Chunk::new(batch.columns().to_vec());
+        writer.write(&chunk, None)?;
         timer.done();
     }
     let timer = disk_write_metric.timer();

ballista/rust/executor/Cargo.toml

@@ -29,8 +29,8 @@ edition = "2018"
 snmalloc = ["snmalloc-rs"]
 
 [dependencies]
-arrow-format = { version = "0.3", features = ["flight-data", "flight-service"] }
-arrow = { package = "arrow2", version="0.8", features = ["io_ipc"] }
+arrow-format = { version = "0.4", features = ["flight-data", "flight-service"] }
+arrow = { package = "arrow2", version="0.9", features = ["io_ipc"] }
 anyhow = "1"
 async-trait = "0.1.36"
 ballista-core = { path = "../core", version = "0.6.0" }

ballista/rust/executor/src/collect.rs

@@ -23,13 +23,12 @@ use std::task::{Context, Poll};
 use std::{any::Any, pin::Pin};
 
 use async_trait::async_trait;
-use datafusion::arrow::{
-    datatypes::SchemaRef, error::Result as ArrowResult, record_batch::RecordBatch,
-};
+use datafusion::arrow::{datatypes::SchemaRef, error::Result as ArrowResult};
 use datafusion::error::DataFusionError;
 use datafusion::physical_plan::{
     DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
 };
+use datafusion::record_batch::RecordBatch;
 use datafusion::{error::Result, physical_plan::RecordBatchStream};
 use futures::stream::SelectAll;
 use futures::Stream;

ballista/rust/executor/src/flight_service.rs

@@ -17,6 +17,8 @@
 
 //! Implementation of the Apache Arrow Flight protocol that wraps an executor.
 
+use arrow::array::ArrayRef;
+use arrow::chunk::Chunk;
 use std::fs::File;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -34,7 +36,6 @@ use arrow_format::flight::data::{
 use arrow_format::flight::service::flight_service_server::FlightService;
 use datafusion::arrow::{
     error::ArrowError, io::ipc::read::FileReader, io::ipc::write::WriteOptions,
-    record_batch::RecordBatch,
 };
 use futures::{Stream, StreamExt};
 use log::{info, warn};
@@ -175,11 +176,11 @@ impl FlightService for BallistaFlightService {
 /// Convert a single RecordBatch into an iterator of FlightData (containing
 /// dictionaries and batches)
 fn create_flight_iter(
-    batch: &RecordBatch,
+    chunk: &Chunk<ArrayRef>,
     options: &WriteOptions,
 ) -> Box<dyn Iterator<Item = Result<FlightData, Status>>> {
     let (flight_dictionaries, flight_batch) =
-        arrow::io::flight::serialize_batch(batch, &[], options);
+        arrow::io::flight::serialize_batch(chunk, &[], options);
     Box::new(
         flight_dictionaries
             .into_iter()
@@ -201,14 +202,13 @@ async fn stream_flight_data(path: String, tx: FlightDataSender) -> Result<(), St
     let reader = FileReader::new(&mut file, file_meta, None);
 
     let options = WriteOptions::default();
-    let schema_flight_data =
-        arrow::io::flight::serialize_schema(reader.schema().as_ref(), &[]);
+    let schema_flight_data = arrow::io::flight::serialize_schema(reader.schema(), None);
     send_response(&tx, Ok(schema_flight_data)).await?;
 
     let mut row_count = 0;
     for batch in reader {
         if let Ok(x) = &batch {
-            row_count += x.num_rows();
+            row_count += x.len();
         }
         let batch_flight_data: Vec<_> = batch
             .map(|b| create_flight_iter(&b, &options).collect())

ballista/rust/scheduler/src/test_utils.rs

@@ -19,6 +19,7 @@ use ballista_core::error::Result;
 
 use datafusion::arrow::datatypes::{DataType, Field, Schema};
 use datafusion::execution::context::{ExecutionConfig, ExecutionContext};
+use datafusion::field_util::SchemaExt;
 use datafusion::prelude::CsvReadOptions;
 
 pub const TPCH_TABLES: &[&str] = &[

benchmarks/.gitignore

@@ -1 +1 @@
-data
+data

benchmarks/Cargo.toml

@@ -32,7 +32,7 @@ simd = ["datafusion/simd"]
 snmalloc = ["snmalloc-rs"]
 
 [dependencies]
-arrow = { package = "arrow2", version="0.8", features = ["io_csv", "io_json", "io_parquet", "io_parquet_compression", "io_ipc", "io_print", "ahash", "compute_merge_sort", "compute", "regex"] }
+arrow = { package = "arrow2", version="0.9", features = ["io_csv", "io_json", "io_parquet", "io_parquet_compression", "io_ipc", "io_print", "ahash", "compute_merge_sort", "compute", "regex"] }
 datafusion = { path = "../datafusion" }
 ballista = { path = "../ballista/rust/client" }
 structopt = { version = "0.3", default-features = false }

benchmarks/src/bin/nyctaxi.rs

@@ -17,6 +17,8 @@
 
 //! Apache Arrow Rust Benchmarks
 
+use arrow::array::ArrayRef;
+use arrow::chunk::Chunk;
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::process;
@@ -28,6 +30,7 @@ use datafusion::arrow::io::print;
 use datafusion::error::Result;
 use datafusion::execution::context::{ExecutionConfig, ExecutionContext};
 
+use datafusion::field_util::SchemaExt;
 use datafusion::physical_plan::collect;
 use datafusion::prelude::CsvReadOptions;
 use structopt::StructOpt;
@@ -124,7 +127,12 @@ async fn execute_sql(ctx: &mut ExecutionContext, sql: &str, debug: bool) -> Resu
     let physical_plan = ctx.create_physical_plan(&plan).await?;
     let result = collect(physical_plan).await?;
     if debug {
-        print::print(&result);
+        let fields = result
+            .first()
+            .map(|b| b.schema().field_names())
+            .unwrap_or(vec![]);
+        let chunks: Vec<Chunk<ArrayRef>> = result.iter().map(|rb| rb.into()).collect();
+        println!("{}", print::write(&chunks, &fields));
     }
     Ok(())
 }