You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ag...@apache.org on 2021/05/30 18:52:00 UTC
[arrow-datafusion] branch master updated: Improved usage of use
arrow in ballista. (#447)
This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/master by this push:
new bdae93b Improved usage of use arrow in ballista. (#447)
bdae93b is described below
commit bdae93b9365ef5892e686915250d42e927d00620
Author: Jorge Leitao <jo...@gmail.com>
AuthorDate: Sun May 30 20:51:52 2021 +0200
Improved usage of use arrow in ballista. (#447)
---
ballista/rust/client/Cargo.toml | 1 -
ballista/rust/client/src/columnar_batch.rs | 2 +-
ballista/rust/client/src/context.rs | 2 +-
ballista/rust/core/Cargo.toml | 1 -
ballista/rust/core/src/client.rs | 12 +-
ballista/rust/core/src/datasource.rs | 2 +-
ballista/rust/core/src/error.rs | 2 +-
.../rust/core/src/execution_plans/query_stage.rs | 2 +-
.../core/src/execution_plans/shuffle_reader.rs | 2 +-
.../core/src/execution_plans/unresolved_shuffle.rs | 2 +-
ballista/rust/core/src/memory_stream.rs | 2 +-
.../rust/core/src/serde/logical_plan/from_proto.rs | 63 +++++----
ballista/rust/core/src/serde/logical_plan/mod.rs | 47 +------
.../rust/core/src/serde/logical_plan/to_proto.rs | 141 +++++++++------------
.../core/src/serde/physical_plan/from_proto.rs | 2 +-
ballista/rust/core/src/serde/physical_plan/mod.rs | 46 +++----
ballista/rust/core/src/serde/scheduler/mod.rs | 4 +-
ballista/rust/core/src/utils.rs | 15 ++-
ballista/rust/executor/src/collect.rs | 6 +-
ballista/rust/executor/src/flight_service.rs | 14 +-
ballista/rust/scheduler/src/test_utils.rs | 2 +-
21 files changed, 157 insertions(+), 213 deletions(-)
diff --git a/ballista/rust/client/Cargo.toml b/ballista/rust/client/Cargo.toml
index 9ce0517..dd1a57c 100644
--- a/ballista/rust/client/Cargo.toml
+++ b/ballista/rust/client/Cargo.toml
@@ -31,5 +31,4 @@ futures = "0.3"
log = "0.4"
tokio = "1.0"
-arrow = { version = "4.0" }
datafusion = { path = "../../../datafusion" }
diff --git a/ballista/rust/client/src/columnar_batch.rs b/ballista/rust/client/src/columnar_batch.rs
index d3ff886..a40b68f 100644
--- a/ballista/rust/client/src/columnar_batch.rs
+++ b/ballista/rust/client/src/columnar_batch.rs
@@ -19,7 +19,7 @@ use std::{collections::HashMap, sync::Arc};
use ballista_core::error::{ballista_error, Result};
-use arrow::{
+use datafusion::arrow::{
array::ArrayRef,
datatypes::{DataType, Schema},
record_batch::RecordBatch,
diff --git a/ballista/rust/client/src/context.rs b/ballista/rust/client/src/context.rs
index df97e3a..4c0ab42 100644
--- a/ballista/rust/client/src/context.rs
+++ b/ballista/rust/client/src/context.rs
@@ -33,7 +33,7 @@ use ballista_core::{
utils::create_datafusion_context,
};
-use arrow::datatypes::Schema;
+use datafusion::arrow::datatypes::Schema;
use datafusion::catalog::TableReference;
use datafusion::error::{DataFusionError, Result};
use datafusion::logical_plan::LogicalPlan;
diff --git a/ballista/rust/core/Cargo.toml b/ballista/rust/core/Cargo.toml
index 2868b60..1fc0e0c 100644
--- a/ballista/rust/core/Cargo.toml
+++ b/ballista/rust/core/Cargo.toml
@@ -40,7 +40,6 @@ tokio = "1.0"
tonic = "0.4"
uuid = { version = "0.8", features = ["v4"] }
-arrow = { version = "4.0" }
arrow-flight = { version = "4.0" }
datafusion = { path = "../../../datafusion" }
diff --git a/ballista/rust/core/src/client.rs b/ballista/rust/core/src/client.rs
index 1d0fedc..c8267c8 100644
--- a/ballista/rust/core/src/client.rs
+++ b/ballista/rust/core/src/client.rs
@@ -31,15 +31,15 @@ use crate::serde::scheduler::{
Action, ExecutePartition, ExecutePartitionResult, PartitionId, PartitionStats,
};
-use arrow::record_batch::RecordBatch;
-use arrow::{
- array::{StringArray, StructArray},
- error::{ArrowError, Result as ArrowResult},
-};
-use arrow::{datatypes::Schema, datatypes::SchemaRef};
use arrow_flight::utils::flight_data_to_arrow_batch;
use arrow_flight::Ticket;
use arrow_flight::{flight_service_client::FlightServiceClient, FlightData};
+use datafusion::arrow::{
+ array::{StringArray, StructArray},
+ datatypes::{Schema, SchemaRef},
+ error::{ArrowError, Result as ArrowResult},
+ record_batch::RecordBatch,
+};
use datafusion::physical_plan::common::collect;
use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream};
use datafusion::{logical_plan::LogicalPlan, physical_plan::RecordBatchStream};
diff --git a/ballista/rust/core/src/datasource.rs b/ballista/rust/core/src/datasource.rs
index 5b1540a..b774b8d 100644
--- a/ballista/rust/core/src/datasource.rs
+++ b/ballista/rust/core/src/datasource.rs
@@ -17,7 +17,7 @@
use std::{any::Any, sync::Arc};
-use arrow::datatypes::SchemaRef;
+use datafusion::arrow::datatypes::SchemaRef;
use datafusion::error::Result as DFResult;
use datafusion::{
datasource::{datasource::Statistics, TableProvider},
diff --git a/ballista/rust/core/src/error.rs b/ballista/rust/core/src/error.rs
index e16920e..b2c8d99 100644
--- a/ballista/rust/core/src/error.rs
+++ b/ballista/rust/core/src/error.rs
@@ -23,7 +23,7 @@ use std::{
io, result,
};
-use arrow::error::ArrowError;
+use datafusion::arrow::error::ArrowError;
use datafusion::error::DataFusionError;
use sqlparser::parser;
diff --git a/ballista/rust/core/src/execution_plans/query_stage.rs b/ballista/rust/core/src/execution_plans/query_stage.rs
index d8822ea..e95a5d8 100644
--- a/ballista/rust/core/src/execution_plans/query_stage.rs
+++ b/ballista/rust/core/src/execution_plans/query_stage.rs
@@ -18,8 +18,8 @@
use std::sync::Arc;
use std::{any::Any, pin::Pin};
-use arrow::datatypes::SchemaRef;
use async_trait::async_trait;
+use datafusion::arrow::datatypes::SchemaRef;
use datafusion::physical_plan::{ExecutionPlan, Partitioning};
use datafusion::{error::Result, physical_plan::RecordBatchStream};
use uuid::Uuid;
diff --git a/ballista/rust/core/src/execution_plans/shuffle_reader.rs b/ballista/rust/core/src/execution_plans/shuffle_reader.rs
index 107cc15..db29cf1 100644
--- a/ballista/rust/core/src/execution_plans/shuffle_reader.rs
+++ b/ballista/rust/core/src/execution_plans/shuffle_reader.rs
@@ -22,8 +22,8 @@ use crate::client::BallistaClient;
use crate::memory_stream::MemoryStream;
use crate::serde::scheduler::PartitionLocation;
-use arrow::datatypes::SchemaRef;
use async_trait::async_trait;
+use datafusion::arrow::datatypes::SchemaRef;
use datafusion::physical_plan::{DisplayFormatType, ExecutionPlan, Partitioning};
use datafusion::{
error::{DataFusionError, Result},
diff --git a/ballista/rust/core/src/execution_plans/unresolved_shuffle.rs b/ballista/rust/core/src/execution_plans/unresolved_shuffle.rs
index 7d147d5..5c1b417 100644
--- a/ballista/rust/core/src/execution_plans/unresolved_shuffle.rs
+++ b/ballista/rust/core/src/execution_plans/unresolved_shuffle.rs
@@ -21,8 +21,8 @@ use std::{any::Any, pin::Pin};
use crate::memory_stream::MemoryStream;
use crate::serde::scheduler::PartitionLocation;
-use arrow::datatypes::SchemaRef;
use async_trait::async_trait;
+use datafusion::arrow::datatypes::SchemaRef;
use datafusion::physical_plan::{ExecutionPlan, Partitioning};
use datafusion::{
error::{DataFusionError, Result},
diff --git a/ballista/rust/core/src/memory_stream.rs b/ballista/rust/core/src/memory_stream.rs
index 8bf5e20..ab72bdc 100644
--- a/ballista/rust/core/src/memory_stream.rs
+++ b/ballista/rust/core/src/memory_stream.rs
@@ -20,7 +20,7 @@
use std::task::{Context, Poll};
-use arrow::{datatypes::SchemaRef, error::Result, record_batch::RecordBatch};
+use datafusion::arrow::{datatypes::SchemaRef, error::Result, record_batch::RecordBatch};
use datafusion::physical_plan::RecordBatchStream;
use futures::Stream;
diff --git a/ballista/rust/core/src/serde/logical_plan/from_proto.rs b/ballista/rust/core/src/serde/logical_plan/from_proto.rs
index 020858f..10c4670 100644
--- a/ballista/rust/core/src/serde/logical_plan/from_proto.rs
+++ b/ballista/rust/core/src/serde/logical_plan/from_proto.rs
@@ -26,7 +26,7 @@ use std::{
unimplemented,
};
-use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use datafusion::logical_plan::{
abs, acos, asin, atan, ceil, cos, exp, floor, ln, log10, log2, round, signum, sin,
sqrt, tan, trunc, Expr, JoinType, LogicalPlan, LogicalPlanBuilder, Operator,
@@ -299,9 +299,9 @@ impl TryInto<datafusion::logical_plan::DFSchemaRef> for protobuf::Schema {
}
}
-impl TryInto<arrow::datatypes::DataType> for &protobuf::scalar_type::Datatype {
+impl TryInto<DataType> for &protobuf::scalar_type::Datatype {
type Error = BallistaError;
- fn try_into(self) -> Result<arrow::datatypes::DataType, Self::Error> {
+ fn try_into(self) -> Result<DataType, Self::Error> {
use protobuf::scalar_type::Datatype;
Ok(match self {
Datatype::Scalar(scalar_type) => {
@@ -332,17 +332,18 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::scalar_type::Datatype {
))
})?;
//Because length is checked above it is safe to unwrap .last()
- let mut scalar_type =
- arrow::datatypes::DataType::List(Box::new(Field::new(
- field_names.last().unwrap().as_str(),
- pb_scalar_type.into(),
- true,
- )));
+ let mut scalar_type = DataType::List(Box::new(Field::new(
+ field_names.last().unwrap().as_str(),
+ pb_scalar_type.into(),
+ true,
+ )));
//Iterate over field names in reverse order except for the last item in the vector
for name in field_names.iter().rev().skip(1) {
- let new_datatype = arrow::datatypes::DataType::List(Box::new(
- Field::new(name.as_str(), scalar_type, true),
- ));
+ let new_datatype = DataType::List(Box::new(Field::new(
+ name.as_str(),
+ scalar_type,
+ true,
+ )));
scalar_type = new_datatype;
}
scalar_type
@@ -351,10 +352,9 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::scalar_type::Datatype {
}
}
-impl TryInto<arrow::datatypes::DataType> for &protobuf::arrow_type::ArrowTypeEnum {
+impl TryInto<DataType> for &protobuf::arrow_type::ArrowTypeEnum {
type Error = BallistaError;
- fn try_into(self) -> Result<arrow::datatypes::DataType, Self::Error> {
- use arrow::datatypes::DataType;
+ fn try_into(self) -> Result<DataType, Self::Error> {
use protobuf::arrow_type;
Ok(match self {
arrow_type::ArrowTypeEnum::None(_) => DataType::Null,
@@ -467,9 +467,8 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::arrow_type::ArrowTypeEnu
}
#[allow(clippy::from_over_into)]
-impl Into<arrow::datatypes::DataType> for protobuf::PrimitiveScalarType {
- fn into(self) -> arrow::datatypes::DataType {
- use arrow::datatypes::DataType;
+impl Into<DataType> for protobuf::PrimitiveScalarType {
+ fn into(self) -> DataType {
match self {
protobuf::PrimitiveScalarType::Bool => DataType::Boolean,
protobuf::PrimitiveScalarType::Uint8 => DataType::UInt8,
@@ -486,10 +485,10 @@ impl Into<arrow::datatypes::DataType> for protobuf::PrimitiveScalarType {
protobuf::PrimitiveScalarType::LargeUtf8 => DataType::LargeUtf8,
protobuf::PrimitiveScalarType::Date32 => DataType::Date32,
protobuf::PrimitiveScalarType::TimeMicrosecond => {
- DataType::Time64(arrow::datatypes::TimeUnit::Microsecond)
+ DataType::Time64(TimeUnit::Microsecond)
}
protobuf::PrimitiveScalarType::TimeNanosecond => {
- DataType::Time64(arrow::datatypes::TimeUnit::Nanosecond)
+ DataType::Time64(TimeUnit::Nanosecond)
}
protobuf::PrimitiveScalarType::Null => DataType::Null,
}
@@ -746,9 +745,9 @@ impl TryInto<datafusion::scalar::ScalarValue> for &protobuf::ScalarListValue {
}
}
-impl TryInto<arrow::datatypes::DataType> for &protobuf::ScalarListType {
+impl TryInto<DataType> for &protobuf::ScalarListType {
type Error = BallistaError;
- fn try_into(self) -> Result<arrow::datatypes::DataType, Self::Error> {
+ fn try_into(self) -> Result<DataType, Self::Error> {
use protobuf::PrimitiveScalarType;
let protobuf::ScalarListType {
deepest_type,
@@ -762,7 +761,7 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::ScalarListType {
));
}
- let mut curr_type = arrow::datatypes::DataType::List(Box::new(Field::new(
+ let mut curr_type = DataType::List(Box::new(Field::new(
//Since checked vector is not empty above this is safe to unwrap
field_names.last().unwrap(),
PrimitiveScalarType::from_i32(*deepest_type)
@@ -774,9 +773,8 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::ScalarListType {
)));
//Iterates over field names in reverse order except for the last item in the vector
for name in field_names.iter().rev().skip(1) {
- let temp_curr_type = arrow::datatypes::DataType::List(Box::new(Field::new(
- name, curr_type, true,
- )));
+ let temp_curr_type =
+ DataType::List(Box::new(Field::new(name, curr_type, true)));
curr_type = temp_curr_type;
}
Ok(curr_type)
@@ -876,8 +874,7 @@ impl TryInto<datafusion::scalar::ScalarValue> for &protobuf::ScalarValue {
.iter()
.map(|val| val.try_into())
.collect::<Result<Vec<_>, _>>()?;
- let scalar_type: arrow::datatypes::DataType =
- pb_scalar_type.try_into()?;
+ let scalar_type: DataType = pb_scalar_type.try_into()?;
ScalarValue::List(Some(typechecked_values), scalar_type)
}
protobuf::scalar_value::Value::NullListValue(v) => {
@@ -1169,9 +1166,9 @@ fn from_proto_binary_op(op: &str) -> Result<Operator, BallistaError> {
}
}
-impl TryInto<arrow::datatypes::DataType> for &protobuf::ScalarType {
+impl TryInto<DataType> for &protobuf::ScalarType {
type Error = BallistaError;
- fn try_into(self) -> Result<arrow::datatypes::DataType, Self::Error> {
+ fn try_into(self) -> Result<DataType, Self::Error> {
let pb_scalartype = self.datatype.as_ref().ok_or_else(|| {
proto_error("ScalarType message missing required field 'datatype'")
})?;
@@ -1202,16 +1199,16 @@ impl TryInto<Schema> for &protobuf::Schema {
}
}
-impl TryInto<arrow::datatypes::Field> for &protobuf::Field {
+impl TryInto<Field> for &protobuf::Field {
type Error = BallistaError;
- fn try_into(self) -> Result<arrow::datatypes::Field, Self::Error> {
+ fn try_into(self) -> Result<Field, Self::Error> {
let pb_datatype = self.arrow_type.as_ref().ok_or_else(|| {
proto_error(
"Protobuf deserialization error: Field message missing required field 'arrow_type'",
)
})?;
- Ok(arrow::datatypes::Field::new(
+ Ok(Field::new(
self.name.as_str(),
pb_datatype.as_ref().try_into()?,
self.nullable,
diff --git a/ballista/rust/core/src/serde/logical_plan/mod.rs b/ballista/rust/core/src/serde/logical_plan/mod.rs
index 48dd96c..d2792b0 100644
--- a/ballista/rust/core/src/serde/logical_plan/mod.rs
+++ b/ballista/rust/core/src/serde/logical_plan/mod.rs
@@ -19,19 +19,18 @@ pub mod from_proto;
pub mod to_proto;
#[cfg(test)]
-
mod roundtrip_tests {
use super::super::{super::error::Result, protobuf};
use crate::error::BallistaError;
- use arrow::datatypes::{DataType, Field, Schema};
use core::panic;
- use datafusion::physical_plan::functions::BuiltinScalarFunction::Sqrt;
use datafusion::{
- logical_plan::{Expr, LogicalPlan, LogicalPlanBuilder},
- physical_plan::csv::CsvReadOptions,
+ arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit},
+ logical_plan::{Expr, LogicalPlan, LogicalPlanBuilder, Partitioning, ToDFSchema},
+ physical_plan::{csv::CsvReadOptions, functions::BuiltinScalarFunction::Sqrt},
prelude::*,
scalar::ScalarValue,
+ sql::parser::FileType,
};
use protobuf::arrow_type;
use std::convert::TryInto;
@@ -57,7 +56,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_repartition() -> Result<()> {
use datafusion::logical_plan::Partitioning;
@@ -119,18 +117,12 @@ mod roundtrip_tests {
Ok(())
}
- fn new_box_field(
- name: &str,
- dt: DataType,
- nullable: bool,
- ) -> Box<arrow::datatypes::Field> {
- Box::new(arrow::datatypes::Field::new(name, dt, nullable))
+ fn new_box_field(name: &str, dt: DataType, nullable: bool) -> Box<Field> {
+ Box::new(Field::new(name, dt, nullable))
}
#[test]
fn scalar_values_error_serialization() -> Result<()> {
- use arrow::datatypes::DataType;
- use datafusion::scalar::ScalarValue;
let should_fail_on_seralize: Vec<ScalarValue> = vec![
//Should fail due to inconsistent types
ScalarValue::List(
@@ -194,8 +186,6 @@ mod roundtrip_tests {
#[test]
fn round_trip_scalar_values() -> Result<()> {
- use arrow::datatypes::DataType;
- use datafusion::scalar::ScalarValue;
let should_pass: Vec<ScalarValue> = vec![
ScalarValue::Boolean(None),
ScalarValue::Float32(None),
@@ -302,8 +292,6 @@ mod roundtrip_tests {
#[test]
fn round_trip_scalar_types() -> Result<()> {
- use arrow::datatypes::DataType;
- use arrow::datatypes::{IntervalUnit, TimeUnit};
let should_pass: Vec<DataType> = vec![
DataType::Boolean,
DataType::Int8,
@@ -459,8 +447,6 @@ mod roundtrip_tests {
#[test]
fn round_trip_datatype() -> Result<()> {
- use arrow::datatypes::DataType;
- use arrow::datatypes::{IntervalUnit, TimeUnit};
let test_cases: Vec<DataType> = vec![
DataType::Null,
DataType::Boolean,
@@ -592,9 +578,6 @@ mod roundtrip_tests {
#[test]
fn roundtrip_null_scalar_values() -> Result<()> {
- use arrow::datatypes::DataType;
- use arrow::datatypes::Field;
- use datafusion::scalar::ScalarValue;
let test_types = vec![
ScalarValue::Boolean(None),
ScalarValue::Float32(None),
@@ -629,7 +612,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_create_external_table() -> Result<()> {
let schema = Schema::new(vec![
Field::new("id", DataType::Int32, false),
@@ -639,12 +621,8 @@ mod roundtrip_tests {
Field::new("salary", DataType::Int32, false),
]);
- use datafusion::logical_plan::ToDFSchema;
-
let df_schema_ref = schema.to_dfschema_ref()?;
- use datafusion::sql::parser::FileType;
-
let filetypes: [FileType; 3] =
[FileType::NdJson, FileType::Parquet, FileType::CSV];
@@ -664,7 +642,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_explain() -> Result<()> {
let schema = Schema::new(vec![
Field::new("id", DataType::Int32, false),
@@ -751,7 +728,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_empty_relation() -> Result<()> {
let plan_false = LogicalPlanBuilder::empty(false)
.build()
@@ -769,7 +745,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_logical_plan() -> Result<()> {
let schema = Schema::new(vec![
Field::new("id", DataType::Int32, false),
@@ -794,7 +769,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_not() -> Result<()> {
let test_expr = Expr::Not(Box::new(Expr::Literal((1.0).into())));
@@ -804,7 +778,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_is_null() -> Result<()> {
let test_expr = Expr::IsNull(Box::new(Expr::Column("id".into())));
@@ -814,7 +787,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_is_not_null() -> Result<()> {
let test_expr = Expr::IsNotNull(Box::new(Expr::Column("id".into())));
@@ -824,7 +796,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_between() -> Result<()> {
let test_expr = Expr::Between {
expr: Box::new(Expr::Literal((1.0).into())),
@@ -839,7 +810,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_case() -> Result<()> {
let test_expr = Expr::Case {
expr: Some(Box::new(Expr::Literal((1.0).into()))),
@@ -856,7 +826,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_cast() -> Result<()> {
let test_expr = Expr::Cast {
expr: Box::new(Expr::Literal((1.0).into())),
@@ -869,7 +838,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_sort_expr() -> Result<()> {
let test_expr = Expr::Sort {
expr: Box::new(Expr::Literal((1.0).into())),
@@ -883,7 +851,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_negative() -> Result<()> {
let test_expr = Expr::Negative(Box::new(Expr::Literal((1.0).into())));
@@ -893,7 +860,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_inlist() -> Result<()> {
let test_expr = Expr::InList {
expr: Box::new(Expr::Literal((1.0).into())),
@@ -907,7 +873,6 @@ mod roundtrip_tests {
}
#[test]
-
fn roundtrip_wildcard() -> Result<()> {
let test_expr = Expr::Wildcard;
diff --git a/ballista/rust/core/src/serde/logical_plan/to_proto.rs b/ballista/rust/core/src/serde/logical_plan/to_proto.rs
index 47e2748..b630dfc 100644
--- a/ballista/rust/core/src/serde/logical_plan/to_proto.rs
+++ b/ballista/rust/core/src/serde/logical_plan/to_proto.rs
@@ -26,7 +26,7 @@ use std::{
use crate::datasource::DfTableAdapter;
use crate::serde::{protobuf, BallistaError};
-use arrow::datatypes::{DataType, Schema};
+use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit};
use datafusion::datasource::CsvFile;
use datafusion::logical_plan::{Expr, JoinType, LogicalPlan};
use datafusion::physical_plan::aggregates::AggregateFunction;
@@ -35,8 +35,8 @@ use datafusion::physical_plan::window_functions::{
};
use datafusion::{datasource::parquet::ParquetTable, logical_plan::exprlist_to_fields};
use protobuf::{
- arrow_type, logical_expr_node::ExprType, scalar_type, DateUnit, Field,
- PrimitiveScalarType, ScalarListValue, ScalarType,
+ arrow_type, logical_expr_node::ExprType, scalar_type, DateUnit, PrimitiveScalarType,
+ ScalarListValue, ScalarType,
};
use sqlparser::ast::{WindowFrame, WindowFrameBound, WindowFrameUnits};
@@ -44,22 +44,17 @@ use super::super::proto_error;
use datafusion::physical_plan::functions::BuiltinScalarFunction;
impl protobuf::IntervalUnit {
- pub fn from_arrow_interval_unit(
- interval_unit: &arrow::datatypes::IntervalUnit,
- ) -> Self {
+ pub fn from_arrow_interval_unit(interval_unit: &IntervalUnit) -> Self {
match interval_unit {
- arrow::datatypes::IntervalUnit::YearMonth => {
- protobuf::IntervalUnit::YearMonth
- }
- arrow::datatypes::IntervalUnit::DayTime => protobuf::IntervalUnit::DayTime,
+ IntervalUnit::YearMonth => protobuf::IntervalUnit::YearMonth,
+ IntervalUnit::DayTime => protobuf::IntervalUnit::DayTime,
}
}
pub fn from_i32_to_arrow(
interval_unit_i32: i32,
- ) -> Result<arrow::datatypes::IntervalUnit, BallistaError> {
+ ) -> Result<IntervalUnit, BallistaError> {
let pb_interval_unit = protobuf::IntervalUnit::from_i32(interval_unit_i32);
- use arrow::datatypes::IntervalUnit;
match pb_interval_unit {
Some(interval_unit) => Ok(match interval_unit {
protobuf::IntervalUnit::YearMonth => IntervalUnit::YearMonth,
@@ -74,15 +69,15 @@ impl protobuf::IntervalUnit {
/* Arrow changed dates to no longer have date unit
impl protobuf::DateUnit {
- pub fn from_arrow_date_unit(val: &arrow::datatypes::DateUnit) -> Self {
+ pub fn from_arrow_date_unit(val: &DateUnit) -> Self {
match val {
- arrow::datatypes::DateUnit::Day => protobuf::DateUnit::Day,
- arrow::datatypes::DateUnit::Millisecond => protobuf::DateUnit::DateMillisecond,
+ DateUnit::Day => protobuf::DateUnit::Day,
+ DateUnit::Millisecond => protobuf::DateUnit::DateMillisecond,
}
}
- pub fn from_i32_to_arrow(date_unit_i32: i32) -> Result<arrow::datatypes::DateUnit, BallistaError> {
+ pub fn from_i32_to_arrow(date_unit_i32: i32) -> Result<DateUnit, BallistaError> {
let pb_date_unit = protobuf::DateUnit::from_i32(date_unit_i32);
- use arrow::datatypes::DateUnit;
+ use datafusion::DateUnit;
match pb_date_unit {
Some(date_unit) => Ok(match date_unit {
protobuf::DateUnit::Day => DateUnit::Day,
@@ -95,21 +90,16 @@ impl protobuf::DateUnit {
}*/
impl protobuf::TimeUnit {
- pub fn from_arrow_time_unit(val: &arrow::datatypes::TimeUnit) -> Self {
+ pub fn from_arrow_time_unit(val: &TimeUnit) -> Self {
match val {
- arrow::datatypes::TimeUnit::Second => protobuf::TimeUnit::Second,
- arrow::datatypes::TimeUnit::Millisecond => {
- protobuf::TimeUnit::TimeMillisecond
- }
- arrow::datatypes::TimeUnit::Microsecond => protobuf::TimeUnit::Microsecond,
- arrow::datatypes::TimeUnit::Nanosecond => protobuf::TimeUnit::Nanosecond,
+ TimeUnit::Second => protobuf::TimeUnit::Second,
+ TimeUnit::Millisecond => protobuf::TimeUnit::TimeMillisecond,
+ TimeUnit::Microsecond => protobuf::TimeUnit::Microsecond,
+ TimeUnit::Nanosecond => protobuf::TimeUnit::Nanosecond,
}
}
- pub fn from_i32_to_arrow(
- time_unit_i32: i32,
- ) -> Result<arrow::datatypes::TimeUnit, BallistaError> {
+ pub fn from_i32_to_arrow(time_unit_i32: i32) -> Result<TimeUnit, BallistaError> {
let pb_time_unit = protobuf::TimeUnit::from_i32(time_unit_i32);
- use arrow::datatypes::TimeUnit;
match pb_time_unit {
Some(time_unit) => Ok(match time_unit {
protobuf::TimeUnit::Second => TimeUnit::Second,
@@ -124,8 +114,8 @@ impl protobuf::TimeUnit {
}
}
-impl From<&arrow::datatypes::Field> for protobuf::Field {
- fn from(field: &arrow::datatypes::Field) -> Self {
+impl From<&Field> for protobuf::Field {
+ fn from(field: &Field) -> Self {
protobuf::Field {
name: field.name().to_owned(),
arrow_type: Some(Box::new(field.data_type().into())),
@@ -135,23 +125,22 @@ impl From<&arrow::datatypes::Field> for protobuf::Field {
}
}
-impl From<&arrow::datatypes::DataType> for protobuf::ArrowType {
- fn from(val: &arrow::datatypes::DataType) -> protobuf::ArrowType {
+impl From<&DataType> for protobuf::ArrowType {
+ fn from(val: &DataType) -> protobuf::ArrowType {
protobuf::ArrowType {
arrow_type_enum: Some(val.into()),
}
}
}
-impl TryInto<arrow::datatypes::DataType> for &protobuf::ArrowType {
+impl TryInto<DataType> for &protobuf::ArrowType {
type Error = BallistaError;
- fn try_into(self) -> Result<arrow::datatypes::DataType, Self::Error> {
+ fn try_into(self) -> Result<DataType, Self::Error> {
let pb_arrow_type = self.arrow_type_enum.as_ref().ok_or_else(|| {
proto_error(
"Protobuf deserialization error: ArrowType missing required field 'data_type'",
)
})?;
- use arrow::datatypes::DataType;
Ok(match pb_arrow_type {
protobuf::arrow_type::ArrowTypeEnum::None(_) => DataType::Null,
protobuf::arrow_type::ArrowTypeEnum::Bool(_) => DataType::Boolean,
@@ -208,7 +197,7 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::ArrowType {
.as_ref()
.ok_or_else(|| proto_error("Protobuf deserialization error: List message was missing required field 'field_type'"))?
.as_ref();
- arrow::datatypes::DataType::List(Box::new(field_ref.try_into()?))
+ DataType::List(Box::new(field_ref.try_into()?))
}
protobuf::arrow_type::ArrowTypeEnum::LargeList(boxed_list) => {
let field_ref = boxed_list
@@ -216,7 +205,7 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::ArrowType {
.as_ref()
.ok_or_else(|| proto_error("Protobuf deserialization error: List message was missing required field 'field_type'"))?
.as_ref();
- arrow::datatypes::DataType::LargeList(Box::new(field_ref.try_into()?))
+ DataType::LargeList(Box::new(field_ref.try_into()?))
}
protobuf::arrow_type::ArrowTypeEnum::FixedSizeList(boxed_list) => {
let fsl_ref = boxed_list.as_ref();
@@ -224,7 +213,7 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::ArrowType {
.field_type
.as_ref()
.ok_or_else(|| proto_error("Protobuf deserialization error: FixedSizeList message was missing required field 'field_type'"))?;
- arrow::datatypes::DataType::FixedSizeList(
+ DataType::FixedSizeList(
Box::new(pb_fieldtype.as_ref().try_into()?),
fsl_ref.list_size,
)
@@ -235,7 +224,7 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::ArrowType {
.iter()
.map(|field| field.try_into())
.collect::<Result<Vec<_>, _>>()?;
- arrow::datatypes::DataType::Struct(fields)
+ DataType::Struct(fields)
}
protobuf::arrow_type::ArrowTypeEnum::Union(union) => {
let union_types = union
@@ -243,7 +232,7 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::ArrowType {
.iter()
.map(|field| field.try_into())
.collect::<Result<Vec<_>, _>>()?;
- arrow::datatypes::DataType::Union(union_types)
+ DataType::Union(union_types)
}
protobuf::arrow_type::ArrowTypeEnum::Dictionary(boxed_dict) => {
let dict_ref = boxed_dict.as_ref();
@@ -255,7 +244,7 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::ArrowType {
.value
.as_ref()
.ok_or_else(|| proto_error("Protobuf deserialization error: Dictionary message was missing required field 'value'"))?;
- arrow::datatypes::DataType::Dictionary(
+ DataType::Dictionary(
Box::new(pb_key.as_ref().try_into()?),
Box::new(pb_value.as_ref().try_into()?),
)
@@ -264,15 +253,15 @@ impl TryInto<arrow::datatypes::DataType> for &protobuf::ArrowType {
}
}
-impl TryInto<arrow::datatypes::DataType> for &Box<protobuf::List> {
+impl TryInto<DataType> for &Box<protobuf::List> {
type Error = BallistaError;
- fn try_into(self) -> Result<arrow::datatypes::DataType, Self::Error> {
+ fn try_into(self) -> Result<DataType, Self::Error> {
let list_ref = self.as_ref();
match &list_ref.field_type {
Some(pb_field) => {
let pb_field_ref = pb_field.as_ref();
- let arrow_field: arrow::datatypes::Field = pb_field_ref.try_into()?;
- Ok(arrow::datatypes::DataType::List(Box::new(arrow_field)))
+ let arrow_field: Field = pb_field_ref.try_into()?;
+ Ok(DataType::List(Box::new(arrow_field)))
}
None => Err(proto_error(
"List message missing required field 'field_type'",
@@ -281,8 +270,8 @@ impl TryInto<arrow::datatypes::DataType> for &Box<protobuf::List> {
}
}
-impl From<&arrow::datatypes::DataType> for protobuf::arrow_type::ArrowTypeEnum {
- fn from(val: &arrow::datatypes::DataType) -> protobuf::arrow_type::ArrowTypeEnum {
+impl From<&DataType> for protobuf::arrow_type::ArrowTypeEnum {
+ fn from(val: &DataType) -> protobuf::arrow_type::ArrowTypeEnum {
use protobuf::arrow_type::ArrowTypeEnum;
use protobuf::ArrowType;
use protobuf::EmptyMessage;
@@ -368,7 +357,7 @@ impl From<&arrow::datatypes::DataType> for protobuf::arrow_type::ArrowTypeEnum {
}
//Does not check if list subtypes are valid
-fn is_valid_scalar_type_no_list_check(datatype: &arrow::datatypes::DataType) -> bool {
+fn is_valid_scalar_type_no_list_check(datatype: &DataType) -> bool {
match datatype {
DataType::Boolean
| DataType::Int8
@@ -384,22 +373,18 @@ fn is_valid_scalar_type_no_list_check(datatype: &arrow::datatypes::DataType) ->
| DataType::LargeUtf8
| DataType::Utf8
| DataType::Date32 => true,
- DataType::Time64(time_unit) => matches!(
- time_unit,
- arrow::datatypes::TimeUnit::Microsecond
- | arrow::datatypes::TimeUnit::Nanosecond
- ),
+ DataType::Time64(time_unit) => {
+ matches!(time_unit, TimeUnit::Microsecond | TimeUnit::Nanosecond)
+ }
DataType::List(_) => true,
_ => false,
}
}
-impl TryFrom<&arrow::datatypes::DataType> for protobuf::scalar_type::Datatype {
+impl TryFrom<&DataType> for protobuf::scalar_type::Datatype {
type Error = BallistaError;
- fn try_from(val: &arrow::datatypes::DataType) -> Result<Self, Self::Error> {
- use protobuf::scalar_type;
- use protobuf::Field;
+ fn try_from(val: &DataType) -> Result<Self, Self::Error> {
use protobuf::{List, PrimitiveScalarType};
let scalar_value = match val {
DataType::Boolean => scalar_type::Datatype::Scalar(PrimitiveScalarType::Bool as i32),
@@ -415,8 +400,8 @@ impl TryFrom<&arrow::datatypes::DataType> for protobuf::scalar_type::Datatype {
DataType::Float64 => scalar_type::Datatype::Scalar(PrimitiveScalarType::Float64 as i32),
DataType::Date32 => scalar_type::Datatype::Scalar(PrimitiveScalarType::Date32 as i32),
DataType::Time64(time_unit) => match time_unit {
- arrow::datatypes::TimeUnit::Microsecond => scalar_type::Datatype::Scalar(PrimitiveScalarType::TimeMicrosecond as i32),
- arrow::datatypes::TimeUnit::Nanosecond => scalar_type::Datatype::Scalar(PrimitiveScalarType::TimeNanosecond as i32),
+ TimeUnit::Microsecond => scalar_type::Datatype::Scalar(PrimitiveScalarType::TimeMicrosecond as i32),
+ TimeUnit::Nanosecond => scalar_type::Datatype::Scalar(PrimitiveScalarType::TimeNanosecond as i32),
_ => {
return Err(proto_error(format!(
"Found invalid time unit for scalar value, only TimeUnit::Microsecond and TimeUnit::Nanosecond are valid time units: {:?}",
@@ -428,7 +413,7 @@ impl TryFrom<&arrow::datatypes::DataType> for protobuf::scalar_type::Datatype {
DataType::LargeUtf8 => scalar_type::Datatype::Scalar(PrimitiveScalarType::LargeUtf8 as i32),
DataType::List(field_type) => {
let mut field_names: Vec<String> = Vec::new();
- let mut curr_field: &arrow::datatypes::Field = field_type.as_ref();
+ let mut curr_field = field_type.as_ref();
field_names.push(curr_field.name().to_owned());
//For each nested field check nested datatype, since datafusion scalars only support recursive lists with a leaf scalar type
// any other compound types are errors.
@@ -458,8 +443,8 @@ impl TryFrom<&arrow::datatypes::DataType> for protobuf::scalar_type::Datatype {
DataType::Float64 => PrimitiveScalarType::Float64,
DataType::Date32 => PrimitiveScalarType::Date32,
DataType::Time64(time_unit) => match time_unit {
- arrow::datatypes::TimeUnit::Microsecond => PrimitiveScalarType::TimeMicrosecond,
- arrow::datatypes::TimeUnit::Nanosecond => PrimitiveScalarType::TimeNanosecond,
+ TimeUnit::Microsecond => PrimitiveScalarType::TimeMicrosecond,
+ TimeUnit::Nanosecond => PrimitiveScalarType::TimeNanosecond,
_ => {
return Err(proto_error(format!(
"Found invalid time unit for scalar value, only TimeUnit::Microsecond and TimeUnit::Nanosecond are valid time units: {:?}",
@@ -594,7 +579,7 @@ impl TryFrom<&datafusion::scalar::ScalarValue> for protobuf::ScalarValue {
let type_checked_values: Vec<protobuf::ScalarValue> = values
.iter()
.map(|scalar| match (scalar, scalar_type) {
- (scalar::ScalarValue::List(_, arrow::datatypes::DataType::List(list_field)), arrow::datatypes::DataType::List(field)) => {
+ (scalar::ScalarValue::List(_, DataType::List(list_field)), DataType::List(field)) => {
let scalar_datatype = field.data_type();
let list_datatype = list_field.data_type();
if std::mem::discriminant(list_datatype) != std::mem::discriminant(scalar_datatype) {
@@ -605,19 +590,19 @@ impl TryFrom<&datafusion::scalar::ScalarValue> for protobuf::ScalarValue {
}
scalar.try_into()
}
- (scalar::ScalarValue::Boolean(_), arrow::datatypes::DataType::Boolean) => scalar.try_into(),
- (scalar::ScalarValue::Float32(_), arrow::datatypes::DataType::Float32) => scalar.try_into(),
- (scalar::ScalarValue::Float64(_), arrow::datatypes::DataType::Float64) => scalar.try_into(),
- (scalar::ScalarValue::Int8(_), arrow::datatypes::DataType::Int8) => scalar.try_into(),
- (scalar::ScalarValue::Int16(_), arrow::datatypes::DataType::Int16) => scalar.try_into(),
- (scalar::ScalarValue::Int32(_), arrow::datatypes::DataType::Int32) => scalar.try_into(),
- (scalar::ScalarValue::Int64(_), arrow::datatypes::DataType::Int64) => scalar.try_into(),
- (scalar::ScalarValue::UInt8(_), arrow::datatypes::DataType::UInt8) => scalar.try_into(),
- (scalar::ScalarValue::UInt16(_), arrow::datatypes::DataType::UInt16) => scalar.try_into(),
- (scalar::ScalarValue::UInt32(_), arrow::datatypes::DataType::UInt32) => scalar.try_into(),
- (scalar::ScalarValue::UInt64(_), arrow::datatypes::DataType::UInt64) => scalar.try_into(),
- (scalar::ScalarValue::Utf8(_), arrow::datatypes::DataType::Utf8) => scalar.try_into(),
- (scalar::ScalarValue::LargeUtf8(_), arrow::datatypes::DataType::LargeUtf8) => scalar.try_into(),
+ (scalar::ScalarValue::Boolean(_), DataType::Boolean) => scalar.try_into(),
+ (scalar::ScalarValue::Float32(_), DataType::Float32) => scalar.try_into(),
+ (scalar::ScalarValue::Float64(_), DataType::Float64) => scalar.try_into(),
+ (scalar::ScalarValue::Int8(_), DataType::Int8) => scalar.try_into(),
+ (scalar::ScalarValue::Int16(_), DataType::Int16) => scalar.try_into(),
+ (scalar::ScalarValue::Int32(_), DataType::Int32) => scalar.try_into(),
+ (scalar::ScalarValue::Int64(_), DataType::Int64) => scalar.try_into(),
+ (scalar::ScalarValue::UInt8(_), DataType::UInt8) => scalar.try_into(),
+ (scalar::ScalarValue::UInt16(_), DataType::UInt16) => scalar.try_into(),
+ (scalar::ScalarValue::UInt32(_), DataType::UInt32) => scalar.try_into(),
+ (scalar::ScalarValue::UInt64(_), DataType::UInt64) => scalar.try_into(),
+ (scalar::ScalarValue::Utf8(_), DataType::Utf8) => scalar.try_into(),
+ (scalar::ScalarValue::LargeUtf8(_), DataType::LargeUtf8) => scalar.try_into(),
_ => Err(proto_error(format!(
"Protobuf serialization error, {:?} was inconsistent with designated type {:?}",
scalar, datatype
@@ -1302,9 +1287,9 @@ impl TryFrom<WindowFrame> for protobuf::WindowFrame {
}
}
-impl TryFrom<&arrow::datatypes::DataType> for protobuf::ScalarType {
+impl TryFrom<&DataType> for protobuf::ScalarType {
type Error = BallistaError;
- fn try_from(value: &arrow::datatypes::DataType) -> Result<Self, Self::Error> {
+ fn try_from(value: &DataType) -> Result<Self, Self::Error> {
let datatype = protobuf::scalar_type::Datatype::try_from(value)?;
Ok(protobuf::ScalarType {
datatype: Some(datatype),
diff --git a/ballista/rust/core/src/serde/physical_plan/from_proto.rs b/ballista/rust/core/src/serde/physical_plan/from_proto.rs
index d034f3c..2039def 100644
--- a/ballista/rust/core/src/serde/physical_plan/from_proto.rs
+++ b/ballista/rust/core/src/serde/physical_plan/from_proto.rs
@@ -28,7 +28,7 @@ use crate::serde::protobuf::LogicalExprNode;
use crate::serde::scheduler::PartitionLocation;
use crate::serde::{proto_error, protobuf};
use crate::{convert_box_required, convert_required};
-use arrow::datatypes::{DataType, Schema, SchemaRef};
+use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
use datafusion::catalog::catalog::{
CatalogList, CatalogProvider, MemoryCatalogList, MemoryCatalogProvider,
};
diff --git a/ballista/rust/core/src/serde/physical_plan/mod.rs b/ballista/rust/core/src/serde/physical_plan/mod.rs
index e7985cc..fdba215 100644
--- a/ballista/rust/core/src/serde/physical_plan/mod.rs
+++ b/ballista/rust/core/src/serde/physical_plan/mod.rs
@@ -20,27 +20,32 @@ pub mod to_proto;
#[cfg(test)]
mod roundtrip_tests {
- use datafusion::physical_plan::hash_utils::JoinType;
use std::{convert::TryInto, sync::Arc};
- use arrow::datatypes::{DataType, Schema};
- use datafusion::physical_plan::ColumnarValue;
- use datafusion::physical_plan::{
- empty::EmptyExec,
- expressions::{Avg, Column, PhysicalSortExpr},
- hash_aggregate::{AggregateMode, HashAggregateExec},
- hash_join::HashJoinExec,
- limit::{GlobalLimitExec, LocalLimitExec},
- sort::SortExec,
- ExecutionPlan,
- };
- use datafusion::physical_plan::{
- AggregateExpr, Distribution, Partitioning, PhysicalExpr,
+ use datafusion::{
+ arrow::{
+ compute::kernels::sort::SortOptions,
+ datatypes::{DataType, Field, Schema},
+ },
+ logical_plan::Operator,
+ physical_plan::{
+ empty::EmptyExec,
+ expressions::{binary, lit, InListExpr, NotExpr},
+ expressions::{Avg, Column, PhysicalSortExpr},
+ filter::FilterExec,
+ hash_aggregate::{AggregateMode, HashAggregateExec},
+ hash_join::{HashJoinExec, PartitionMode},
+ hash_utils::JoinType,
+ limit::{GlobalLimitExec, LocalLimitExec},
+ sort::SortExec,
+ AggregateExpr, ColumnarValue, Distribution, ExecutionPlan, Partitioning,
+ PhysicalExpr,
+ },
+ scalar::ScalarValue,
};
use super::super::super::error::Result;
use super::super::protobuf;
- use datafusion::physical_plan::hash_join::PartitionMode;
fn roundtrip_test(exec_plan: Arc<dyn ExecutionPlan>) -> Result<()> {
let proto: protobuf::PhysicalPlanNode = exec_plan.clone().try_into()?;
@@ -75,7 +80,6 @@ mod roundtrip_tests {
#[test]
fn roundtrip_hash_join() -> Result<()> {
- use arrow::datatypes::{DataType, Field, Schema};
let field_a = Field::new("col", DataType::Int64, false);
let schema_left = Schema::new(vec![field_a.clone()]);
let schema_right = Schema::new(vec![field_a]);
@@ -95,7 +99,6 @@ mod roundtrip_tests {
#[test]
fn rountrip_hash_aggregate() -> Result<()> {
- use arrow::datatypes::{DataType, Field, Schema};
let groups: Vec<(Arc<dyn PhysicalExpr>, String)> =
vec![(col("a"), "unused".to_string())];
@@ -120,13 +123,6 @@ mod roundtrip_tests {
#[test]
fn roundtrip_filter_with_not_and_in_list() -> Result<()> {
- use arrow::datatypes::{DataType, Field, Schema};
- use datafusion::logical_plan::Operator;
- use datafusion::physical_plan::{
- expressions::{binary, lit, InListExpr, NotExpr},
- filter::FilterExec,
- };
- use datafusion::scalar::ScalarValue;
let field_a = Field::new("a", DataType::Boolean, false);
let field_b = Field::new("b", DataType::Int64, false);
let field_c = Field::new("c", DataType::Int64, false);
@@ -149,8 +145,6 @@ mod roundtrip_tests {
#[test]
fn roundtrip_sort() -> Result<()> {
- use arrow::compute::kernels::sort::SortOptions;
- use arrow::datatypes::{DataType, Field, Schema};
let field_a = Field::new("a", DataType::Boolean, false);
let field_b = Field::new("b", DataType::Int64, false);
let schema = Arc::new(Schema::new(vec![field_a, field_b]));
diff --git a/ballista/rust/core/src/serde/scheduler/mod.rs b/ballista/rust/core/src/serde/scheduler/mod.rs
index bbbd48b..b502c32 100644
--- a/ballista/rust/core/src/serde/scheduler/mod.rs
+++ b/ballista/rust/core/src/serde/scheduler/mod.rs
@@ -17,10 +17,10 @@
use std::{collections::HashMap, sync::Arc};
-use arrow::array::{
+use datafusion::arrow::array::{
ArrayBuilder, ArrayRef, StructArray, StructBuilder, UInt64Array, UInt64Builder,
};
-use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion::logical_plan::LogicalPlan;
use datafusion::physical_plan::ExecutionPlan;
use serde::Serialize;
diff --git a/ballista/rust/core/src/utils.rs b/ballista/rust/core/src/utils.rs
index dc570f8..85d1161 100644
--- a/ballista/rust/core/src/utils.rs
+++ b/ballista/rust/core/src/utils.rs
@@ -26,13 +26,16 @@ use crate::error::{BallistaError, Result};
use crate::execution_plans::{QueryStageExec, UnresolvedShuffleExec};
use crate::memory_stream::MemoryStream;
use crate::serde::scheduler::PartitionStats;
-use arrow::array::{
- ArrayBuilder, ArrayRef, StructArray, StructBuilder, UInt64Array, UInt64Builder,
+
+use datafusion::arrow::{
+ array::{
+ ArrayBuilder, ArrayRef, StructArray, StructBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field},
+ ipc::reader::FileReader,
+ ipc::writer::FileWriter,
+ record_batch::RecordBatch,
};
-use arrow::datatypes::{DataType, Field};
-use arrow::ipc::reader::FileReader;
-use arrow::ipc::writer::FileWriter;
-use arrow::record_batch::RecordBatch;
use datafusion::execution::context::{ExecutionConfig, ExecutionContext};
use datafusion::logical_plan::Operator;
use datafusion::physical_optimizer::coalesce_batches::CoalesceBatches;
diff --git a/ballista/rust/executor/src/collect.rs b/ballista/rust/executor/src/collect.rs
index a2f9d4c..a4c544f 100644
--- a/ballista/rust/executor/src/collect.rs
+++ b/ballista/rust/executor/src/collect.rs
@@ -22,10 +22,10 @@ use std::sync::Arc;
use std::task::{Context, Poll};
use std::{any::Any, pin::Pin};
-use arrow::datatypes::SchemaRef;
-use arrow::error::Result as ArrowResult;
-use arrow::record_batch::RecordBatch;
use async_trait::async_trait;
+use datafusion::arrow::{
+ datatypes::SchemaRef, error::Result as ArrowResult, record_batch::RecordBatch,
+};
use datafusion::error::DataFusionError;
use datafusion::physical_plan::{ExecutionPlan, Partitioning, SendableRecordBatchStream};
use datafusion::{error::Result, physical_plan::RecordBatchStream};
diff --git a/ballista/rust/executor/src/flight_service.rs b/ballista/rust/executor/src/flight_service.rs
index 62aaf7f..b35ac15 100644
--- a/ballista/rust/executor/src/flight_service.rs
+++ b/ballista/rust/executor/src/flight_service.rs
@@ -28,17 +28,19 @@ use ballista_core::serde::decode_protobuf;
use ballista_core::serde::scheduler::{Action as BallistaAction, PartitionStats};
use ballista_core::utils;
-use arrow::array::{ArrayRef, StringBuilder};
-use arrow::datatypes::{DataType, Field, Schema};
-use arrow::error::ArrowError;
-use arrow::ipc::reader::FileReader;
-use arrow::ipc::writer::IpcWriteOptions;
-use arrow::record_batch::RecordBatch;
use arrow_flight::{
flight_service_server::FlightService, Action, ActionType, Criteria, Empty,
FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse,
PutResult, SchemaResult, Ticket,
};
+use datafusion::arrow::{
+ array::{ArrayRef, StringBuilder},
+ datatypes::{DataType, Field, Schema},
+ error::ArrowError,
+ ipc::reader::FileReader,
+ ipc::writer::IpcWriteOptions,
+ record_batch::RecordBatch,
+};
use datafusion::error::DataFusionError;
use datafusion::physical_plan::displayable;
use futures::{Stream, StreamExt};
diff --git a/ballista/rust/scheduler/src/test_utils.rs b/ballista/rust/scheduler/src/test_utils.rs
index 0989060..311f9a7 100644
--- a/ballista/rust/scheduler/src/test_utils.rs
+++ b/ballista/rust/scheduler/src/test_utils.rs
@@ -19,7 +19,7 @@ use std::sync::Arc;
use ballista_core::error::Result;
-use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::execution::context::{ExecutionConfig, ExecutionContext};
use datafusion::physical_optimizer::coalesce_batches::CoalesceBatches;
use datafusion::physical_optimizer::merge_exec::AddMergeExec;