You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2022/02/09 18:59:38 UTC
[arrow-datafusion] branch master updated: Abstract over logical and physical plan representations in Ballista (#1677)
This is an automated email from the ASF dual-hosted git repository.
alamb 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 1bedaf3 Abstract over logical and physical plan representations in Ballista (#1677)
1bedaf3 is described below
commit 1bedaf33986ea01d78d21cf1878da40497b1b828
Author: Dan Harris <13...@users.noreply.github.com>
AuthorDate: Wed Feb 9 13:59:32 2022 -0500
Abstract over logical and physical plan representations in Ballista (#1677)
* Abstract over logical plan representation in Ballista
* Linter fix
* Create trait for ExecutionPlan serializable representation and make the ExecutionContext global to the server
* Make FileScanConfig decoding context aware
* Clippy fix
* More clippy issues and use global context everywehere
* Fix various cases in which we were still using the old codec
* Remove unused code and add test case for logical plan round trip with custom object store
* clippy is a cruel tyrant
* Decode execution plan from raw bytes in protobuf message
* clippy fixes
* Remove commented code
* Add LogicalExtensionCodec design
* Implement extension codecs for serializing user-defined logical and physical plans
* appease clippy
* remove commented code
---
ballista/rust/client/src/context.rs | 9 +-
ballista/rust/core/proto/ballista.proto | 16 +-
.../core/src/execution_plans/distributed_query.rs | 79 +-
.../rust/core/src/serde/logical_plan/from_proto.rs | 302 -------
ballista/rust/core/src/serde/logical_plan/mod.rs | 966 ++++++++++++++++++++-
.../rust/core/src/serde/logical_plan/to_proto.rs | 349 --------
ballista/rust/core/src/serde/mod.rs | 579 +++++++++++-
.../core/src/serde/physical_plan/from_proto.rs | 382 --------
ballista/rust/core/src/serde/physical_plan/mod.rs | 923 +++++++++++++++++++-
.../rust/core/src/serde/physical_plan/to_proto.rs | 333 -------
ballista/rust/core/src/serde/scheduler/to_proto.rs | 17 -
ballista/rust/core/src/utils.rs | 52 +-
ballista/rust/executor/src/execution_loop.rs | 19 +-
ballista/rust/executor/src/executor.rs | 10 +-
ballista/rust/executor/src/executor_server.rs | 44 +-
ballista/rust/executor/src/main.rs | 12 +-
ballista/rust/executor/src/standalone.rs | 13 +-
ballista/rust/scheduler/src/api/handlers.rs | 5 +-
ballista/rust/scheduler/src/api/mod.rs | 12 +-
ballista/rust/scheduler/src/lib.rs | 142 ++-
ballista/rust/scheduler/src/main.rs | 46 +-
ballista/rust/scheduler/src/planner.rs | 18 +-
ballista/rust/scheduler/src/standalone.rs | 16 +-
ballista/rust/scheduler/src/state/mod.rs | 166 ++--
benchmarks/src/bin/tpch.rs | 41 +-
25 files changed, 2971 insertions(+), 1580 deletions(-)
diff --git a/ballista/rust/client/src/context.rs b/ballista/rust/client/src/context.rs
index 4cd5a21..8326afc 100644
--- a/ballista/rust/client/src/context.rs
+++ b/ballista/rust/client/src/context.rs
@@ -25,6 +25,7 @@ use std::path::PathBuf;
use std::sync::Arc;
use ballista_core::config::BallistaConfig;
+use ballista_core::serde::protobuf::LogicalPlanNode;
use ballista_core::utils::create_df_ctx_with_ballista_query_planner;
use datafusion::catalog::TableReference;
@@ -144,7 +145,7 @@ impl BallistaContext {
// use local DataFusion context for now but later this might call the scheduler
let mut ctx = {
let guard = self.state.lock();
- create_df_ctx_with_ballista_query_planner(
+ create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
&guard.scheduler_host,
guard.scheduler_port,
guard.config(),
@@ -164,7 +165,7 @@ impl BallistaContext {
// use local DataFusion context for now but later this might call the scheduler
let mut ctx = {
let guard = self.state.lock();
- create_df_ctx_with_ballista_query_planner(
+ create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
&guard.scheduler_host,
guard.scheduler_port,
guard.config(),
@@ -188,7 +189,7 @@ impl BallistaContext {
// use local DataFusion context for now but later this might call the scheduler
let mut ctx = {
let guard = self.state.lock();
- create_df_ctx_with_ballista_query_planner(
+ create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
&guard.scheduler_host,
guard.scheduler_port,
guard.config(),
@@ -282,7 +283,7 @@ impl BallistaContext {
pub async fn sql(&self, sql: &str) -> Result<Arc<dyn DataFrame>> {
let mut ctx = {
let state = self.state.lock();
- create_df_ctx_with_ballista_query_planner(
+ create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
&state.scheduler_host,
state.scheduler_port,
state.config(),
diff --git a/ballista/rust/core/proto/ballista.proto b/ballista/rust/core/proto/ballista.proto
index 0b0d364..32e363d 100644
--- a/ballista/rust/core/proto/ballista.proto
+++ b/ballista/rust/core/proto/ballista.proto
@@ -266,9 +266,15 @@ message LogicalPlanNode {
AnalyzeNode analyze = 14;
CrossJoinNode cross_join = 15;
ValuesNode values = 16;
+ LogicalExtensionNode extension = 17;
}
}
+message LogicalExtensionNode {
+ bytes node = 1;
+ repeated LogicalPlanNode inputs = 2;
+}
+
message ProjectionColumns {
repeated string columns = 1;
}
@@ -488,9 +494,15 @@ message PhysicalPlanNode {
ShuffleWriterExecNode shuffle_writer = 18;
CrossJoinExecNode cross_join = 19;
AvroScanExecNode avro_scan = 20;
+ PhysicalExtensionNode extension = 21;
}
}
+message PhysicalExtensionNode {
+ bytes node = 1;
+ repeated PhysicalPlanNode inputs = 2;
+}
+
// physical expressions
message PhysicalExprNode {
oneof ExprType {
@@ -934,7 +946,7 @@ message PollWorkParams {
message TaskDefinition {
PartitionId task_id = 1;
- PhysicalPlanNode plan = 2;
+ bytes plan = 2;
// Output partition for shuffle writer
PhysicalHashRepartition output_partitioning = 3;
}
@@ -980,7 +992,7 @@ message UpdateTaskStatusResult {
message ExecuteQueryParams {
oneof query {
- LogicalPlanNode logical_plan = 1;
+ bytes logical_plan = 1;
string sql = 2;
}
repeated KeyValuePair settings = 3;
diff --git a/ballista/rust/core/src/execution_plans/distributed_query.rs b/ballista/rust/core/src/execution_plans/distributed_query.rs
index 2216015..d6b3c3d 100644
--- a/ballista/rust/core/src/execution_plans/distributed_query.rs
+++ b/ballista/rust/core/src/execution_plans/distributed_query.rs
@@ -17,6 +17,9 @@
use std::any::Any;
use std::convert::TryInto;
+use std::fmt::Debug;
+use std::marker::PhantomData;
+use std::marker::Send;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
@@ -39,6 +42,7 @@ use datafusion::physical_plan::{
SendableRecordBatchStream, Statistics,
};
+use crate::serde::{AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec};
use async_trait::async_trait;
use datafusion::execution::runtime_env::RuntimeEnv;
use futures::future;
@@ -50,27 +54,64 @@ use log::{error, info};
/// batches directly from the executors that hold the results from the final
/// query stage.
#[derive(Debug, Clone)]
-pub struct DistributedQueryExec {
+pub struct DistributedQueryExec<T: 'static + AsLogicalPlan> {
/// Ballista scheduler URL
scheduler_url: String,
/// Ballista configuration
config: BallistaConfig,
/// Logical plan to execute
plan: LogicalPlan,
+ /// Codec for LogicalPlan extensions
+ extension_codec: Arc<dyn LogicalExtensionCodec>,
+ /// Phantom data for serializable plan message
+ plan_repr: PhantomData<T>,
}
-impl DistributedQueryExec {
+impl<T: 'static + AsLogicalPlan> DistributedQueryExec<T> {
pub fn new(scheduler_url: String, config: BallistaConfig, plan: LogicalPlan) -> Self {
Self {
scheduler_url,
config,
plan,
+ extension_codec: Arc::new(DefaultLogicalExtensionCodec {}),
+ plan_repr: PhantomData,
+ }
+ }
+
+ pub fn with_extension(
+ scheduler_url: String,
+ config: BallistaConfig,
+ plan: LogicalPlan,
+ extension_codec: Arc<dyn LogicalExtensionCodec>,
+ ) -> Self {
+ Self {
+ scheduler_url,
+ config,
+ plan,
+ extension_codec,
+ plan_repr: PhantomData,
+ }
+ }
+
+ pub fn with_repr(
+ scheduler_url: String,
+ config: BallistaConfig,
+ plan: LogicalPlan,
+ extension_codec: Arc<dyn LogicalExtensionCodec>,
+ plan_repr: PhantomData<T>,
+ ) -> Self {
+ Self {
+ scheduler_url,
+ config,
+ plan,
+ extension_codec,
+ plan_repr,
}
}
}
#[async_trait]
-impl ExecutionPlan for DistributedQueryExec {
+impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
fn as_any(&self) -> &dyn Any {
self
}
@@ -99,11 +140,13 @@ impl ExecutionPlan for DistributedQueryExec {
&self,
_children: Vec<Arc<dyn ExecutionPlan>>,
) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
- Ok(Arc::new(DistributedQueryExec::new(
- self.scheduler_url.clone(),
- self.config.clone(),
- self.plan.clone(),
- )))
+ Ok(Arc::new(DistributedQueryExec {
+ scheduler_url: self.scheduler_url.clone(),
+ config: self.config.clone(),
+ plan: self.plan.clone(),
+ extension_codec: self.extension_codec.clone(),
+ plan_repr: self.plan_repr,
+ }))
}
async fn execute(
@@ -121,13 +164,23 @@ impl ExecutionPlan for DistributedQueryExec {
let schema: Schema = self.plan.schema().as_ref().clone().into();
+ let mut buf: Vec<u8> = vec![];
+ let plan_message =
+ T::try_from_logical_plan(&self.plan, self.extension_codec.as_ref()).map_err(
+ |e| {
+ DataFusionError::Internal(format!(
+ "failed to serialize logical plan: {:?}",
+ e
+ ))
+ },
+ )?;
+ plan_message.try_encode(&mut buf).map_err(|e| {
+ DataFusionError::Execution(format!("failed to encode logical plan: {:?}", e))
+ })?;
+
let job_id = scheduler
.execute_query(ExecuteQueryParams {
- query: Some(Query::LogicalPlan(
- (&self.plan)
- .try_into()
- .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?,
- )),
+ query: Some(Query::LogicalPlan(buf)),
settings: self
.config
.settings()
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 044f823..5951376 100644
--- a/ballista/rust/core/src/serde/logical_plan/from_proto.rs
+++ b/ballista/rust/core/src/serde/logical_plan/from_proto.rs
@@ -51,308 +51,6 @@ use std::{
unimplemented,
};
-impl TryInto<LogicalPlan> for &protobuf::LogicalPlanNode {
- type Error = BallistaError;
-
- fn try_into(self) -> Result<LogicalPlan, Self::Error> {
- let plan = self.logical_plan_type.as_ref().ok_or_else(|| {
- proto_error(format!(
- "logical_plan::from_proto() Unsupported logical plan '{:?}'",
- self
- ))
- })?;
- match plan {
- LogicalPlanType::Values(values) => {
- let n_cols = values.n_cols as usize;
- let values: Vec<Vec<Expr>> = if values.values_list.is_empty() {
- Ok(Vec::new())
- } else if values.values_list.len() % n_cols != 0 {
- Err(BallistaError::General(format!(
- "Invalid values list length, expect {} to be divisible by {}",
- values.values_list.len(),
- n_cols
- )))
- } else {
- values
- .values_list
- .chunks_exact(n_cols)
- .map(|r| {
- r.iter()
- .map(|v| v.try_into())
- .collect::<Result<Vec<_>, _>>()
- })
- .collect::<Result<Vec<_>, _>>()
- }?;
- LogicalPlanBuilder::values(values)?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::Projection(projection) => {
- let input: LogicalPlan = convert_box_required!(projection.input)?;
- let x: Vec<Expr> = projection
- .expr
- .iter()
- .map(|expr| expr.try_into())
- .collect::<Result<Vec<_>, _>>()?;
- LogicalPlanBuilder::from(input)
- .project_with_alias(
- x,
- projection.optional_alias.as_ref().map(|a| match a {
- protobuf::projection_node::OptionalAlias::Alias(alias) => {
- alias.clone()
- }
- }),
- )?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::Selection(selection) => {
- let input: LogicalPlan = convert_box_required!(selection.input)?;
- let expr: Expr = selection
- .expr
- .as_ref()
- .ok_or_else(|| {
- BallistaError::General("expression required".to_string())
- })?
- .try_into()?;
- LogicalPlanBuilder::from(input)
- .filter(expr)?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::Window(window) => {
- let input: LogicalPlan = convert_box_required!(window.input)?;
- let window_expr = window
- .window_expr
- .iter()
- .map(|expr| expr.try_into())
- .collect::<Result<Vec<Expr>, _>>()?;
- LogicalPlanBuilder::from(input)
- .window(window_expr)?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::Aggregate(aggregate) => {
- let input: LogicalPlan = convert_box_required!(aggregate.input)?;
- let group_expr = aggregate
- .group_expr
- .iter()
- .map(|expr| expr.try_into())
- .collect::<Result<Vec<Expr>, _>>()?;
- let aggr_expr = aggregate
- .aggr_expr
- .iter()
- .map(|expr| expr.try_into())
- .collect::<Result<Vec<Expr>, _>>()?;
- LogicalPlanBuilder::from(input)
- .aggregate(group_expr, aggr_expr)?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::ListingScan(scan) => {
- let schema: Schema = convert_required!(scan.schema)?;
-
- let mut projection = None;
- if let Some(columns) = &scan.projection {
- let column_indices = columns
- .columns
- .iter()
- .map(|name| schema.index_of(name))
- .collect::<Result<Vec<usize>, _>>()?;
- projection = Some(column_indices);
- }
-
- let filters = scan
- .filters
- .iter()
- .map(|e| e.try_into())
- .collect::<Result<Vec<_>, _>>()?;
-
- let file_format: Arc<dyn FileFormat> =
- match scan.file_format_type.as_ref().ok_or_else(|| {
- proto_error(format!(
- "logical_plan::from_proto() Unsupported file format '{:?}'",
- self
- ))
- })? {
- &FileFormatType::Parquet(protobuf::ParquetFormat {
- enable_pruning,
- }) => Arc::new(
- ParquetFormat::default().with_enable_pruning(enable_pruning),
- ),
- FileFormatType::Csv(protobuf::CsvFormat {
- has_header,
- delimiter,
- }) => Arc::new(
- CsvFormat::default()
- .with_has_header(*has_header)
- .with_delimiter(str_to_byte(delimiter)?),
- ),
- FileFormatType::Avro(..) => Arc::new(AvroFormat::default()),
- };
-
- let options = ListingOptions {
- file_extension: scan.file_extension.clone(),
- format: file_format,
- table_partition_cols: scan.table_partition_cols.clone(),
- collect_stat: scan.collect_stat,
- target_partitions: scan.target_partitions as usize,
- };
-
- let provider = ListingTable::new(
- Arc::new(LocalFileSystem {}),
- scan.path.clone(),
- Arc::new(schema),
- options,
- );
-
- LogicalPlanBuilder::scan_with_filters(
- &scan.table_name,
- Arc::new(provider),
- projection,
- filters,
- )?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::Sort(sort) => {
- let input: LogicalPlan = convert_box_required!(sort.input)?;
- let sort_expr: Vec<Expr> = sort
- .expr
- .iter()
- .map(|expr| expr.try_into())
- .collect::<Result<Vec<Expr>, _>>()?;
- LogicalPlanBuilder::from(input)
- .sort(sort_expr)?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::Repartition(repartition) => {
- use datafusion::logical_plan::Partitioning;
- let input: LogicalPlan = convert_box_required!(repartition.input)?;
- use protobuf::repartition_node::PartitionMethod;
- let pb_partition_method = repartition.partition_method.clone().ok_or_else(|| {
- BallistaError::General(String::from(
- "Protobuf deserialization error, RepartitionNode was missing required field 'partition_method'",
- ))
- })?;
-
- let partitioning_scheme = match pb_partition_method {
- PartitionMethod::Hash(protobuf::HashRepartition {
- hash_expr: pb_hash_expr,
- partition_count,
- }) => Partitioning::Hash(
- pb_hash_expr
- .iter()
- .map(|pb_expr| pb_expr.try_into())
- .collect::<Result<Vec<_>, _>>()?,
- partition_count as usize,
- ),
- PartitionMethod::RoundRobin(partition_count) => {
- Partitioning::RoundRobinBatch(partition_count as usize)
- }
- };
-
- LogicalPlanBuilder::from(input)
- .repartition(partitioning_scheme)?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::EmptyRelation(empty_relation) => {
- LogicalPlanBuilder::empty(empty_relation.produce_one_row)
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::CreateExternalTable(create_extern_table) => {
- let pb_schema = (create_extern_table.schema.clone()).ok_or_else(|| {
- BallistaError::General(String::from(
- "Protobuf deserialization error, CreateExternalTableNode was missing required field schema.",
- ))
- })?;
-
- let pb_file_type: protobuf::FileType =
- create_extern_table.file_type.try_into()?;
-
- Ok(LogicalPlan::CreateExternalTable(CreateExternalTable {
- schema: pb_schema.try_into()?,
- name: create_extern_table.name.clone(),
- location: create_extern_table.location.clone(),
- file_type: pb_file_type.into(),
- has_header: create_extern_table.has_header,
- }))
- }
- LogicalPlanType::Analyze(analyze) => {
- let input: LogicalPlan = convert_box_required!(analyze.input)?;
- LogicalPlanBuilder::from(input)
- .explain(analyze.verbose, true)?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::Explain(explain) => {
- let input: LogicalPlan = convert_box_required!(explain.input)?;
- LogicalPlanBuilder::from(input)
- .explain(explain.verbose, false)?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::Limit(limit) => {
- let input: LogicalPlan = convert_box_required!(limit.input)?;
- LogicalPlanBuilder::from(input)
- .limit(limit.limit as usize)?
- .build()
- .map_err(|e| e.into())
- }
- LogicalPlanType::Join(join) => {
- let left_keys: Vec<Column> =
- join.left_join_column.iter().map(|i| i.into()).collect();
- let right_keys: Vec<Column> =
- join.right_join_column.iter().map(|i| i.into()).collect();
- let join_type =
- protobuf::JoinType::from_i32(join.join_type).ok_or_else(|| {
- proto_error(format!(
- "Received a JoinNode message with unknown JoinType {}",
- join.join_type
- ))
- })?;
- let join_constraint = protobuf::JoinConstraint::from_i32(
- join.join_constraint,
- )
- .ok_or_else(|| {
- proto_error(format!(
- "Received a JoinNode message with unknown JoinConstraint {}",
- join.join_constraint
- ))
- })?;
-
- let builder = LogicalPlanBuilder::from(convert_box_required!(join.left)?);
- let builder = match join_constraint.into() {
- JoinConstraint::On => builder.join(
- &convert_box_required!(join.right)?,
- join_type.into(),
- (left_keys, right_keys),
- )?,
- JoinConstraint::Using => builder.join_using(
- &convert_box_required!(join.right)?,
- join_type.into(),
- left_keys,
- )?,
- };
-
- builder.build().map_err(|e| e.into())
- }
- LogicalPlanType::CrossJoin(crossjoin) => {
- let left = convert_box_required!(crossjoin.left)?;
- let right = convert_box_required!(crossjoin.right)?;
-
- LogicalPlanBuilder::from(left)
- .cross_join(&right)?
- .build()
- .map_err(|e| e.into())
- }
- }
- }
-}
-
impl From<&protobuf::Column> for Column {
fn from(c: &protobuf::Column) -> Column {
let c = c.clone();
diff --git a/ballista/rust/core/src/serde/logical_plan/mod.rs b/ballista/rust/core/src/serde/logical_plan/mod.rs
index c00e3e4..31c3308 100644
--- a/ballista/rust/core/src/serde/logical_plan/mod.rs
+++ b/ballista/rust/core/src/serde/logical_plan/mod.rs
@@ -15,15 +15,849 @@
// specific language governing permissions and limitations
// under the License.
+use crate::error::BallistaError;
+use crate::serde::protobuf::LogicalExtensionNode;
+use crate::serde::{
+ byte_to_string, proto_error, protobuf, str_to_byte, AsLogicalPlan,
+ LogicalExtensionCodec,
+};
+use crate::{convert_required, into_logical_plan};
+use datafusion::arrow::datatypes::Schema;
+use datafusion::datasource::file_format::avro::AvroFormat;
+use datafusion::datasource::file_format::csv::CsvFormat;
+use datafusion::datasource::file_format::parquet::ParquetFormat;
+use datafusion::datasource::file_format::FileFormat;
+use datafusion::datasource::listing::{ListingOptions, ListingTable};
+use datafusion::datasource::object_store::local::LocalFileSystem;
+use datafusion::logical_plan::plan::Extension;
+use datafusion::logical_plan::plan::{
+ Aggregate, EmptyRelation, Filter, Join, Projection, Sort, Window,
+};
+use datafusion::logical_plan::{
+ exprlist_to_fields,
+ window_frames::{WindowFrame, WindowFrameBound, WindowFrameUnits},
+ Column, CreateExternalTable, CrossJoin, Expr, JoinConstraint, JoinType, Limit,
+ LogicalPlan, LogicalPlanBuilder, Repartition, TableScan, Values,
+};
+use datafusion::prelude::ExecutionContext;
+use log::error;
+use prost::bytes::{Buf, BufMut};
+use prost::Message;
+use protobuf::listing_table_scan_node::FileFormatType;
+use protobuf::logical_plan_node::LogicalPlanType;
+use protobuf::{logical_expr_node::ExprType, scalar_type, LogicalPlanNode};
+use std::convert::TryInto;
+use std::sync::Arc;
+
pub mod from_proto;
pub mod to_proto;
+impl AsLogicalPlan for LogicalPlanNode {
+ fn try_decode(buf: &[u8]) -> Result<Self, BallistaError>
+ where
+ Self: Sized,
+ {
+ LogicalPlanNode::decode(buf).map_err(|e| {
+ BallistaError::Internal(format!("failed to decode logical plan: {:?}", e))
+ })
+ }
+
+ fn try_encode<B>(&self, buf: &mut B) -> Result<(), BallistaError>
+ where
+ B: BufMut,
+ Self: Sized,
+ {
+ self.encode(buf).map_err(|e| {
+ BallistaError::Internal(format!("failed to encode logical plan: {:?}", e))
+ })
+ }
+
+ fn try_into_logical_plan(
+ &self,
+ ctx: &ExecutionContext,
+ extension_codec: &dyn LogicalExtensionCodec,
+ ) -> Result<LogicalPlan, BallistaError> {
+ let plan = self.logical_plan_type.as_ref().ok_or_else(|| {
+ proto_error(format!(
+ "logical_plan::from_proto() Unsupported logical plan '{:?}'",
+ self
+ ))
+ })?;
+ match plan {
+ LogicalPlanType::Values(values) => {
+ let n_cols = values.n_cols as usize;
+ let values: Vec<Vec<Expr>> = if values.values_list.is_empty() {
+ Ok(Vec::new())
+ } else if values.values_list.len() % n_cols != 0 {
+ Err(BallistaError::General(format!(
+ "Invalid values list length, expect {} to be divisible by {}",
+ values.values_list.len(),
+ n_cols
+ )))
+ } else {
+ values
+ .values_list
+ .chunks_exact(n_cols)
+ .map(|r| {
+ r.iter()
+ .map(|v| v.try_into())
+ .collect::<Result<Vec<_>, _>>()
+ })
+ .collect::<Result<Vec<_>, _>>()
+ }?;
+ LogicalPlanBuilder::values(values)?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::Projection(projection) => {
+ let input: LogicalPlan =
+ into_logical_plan!(projection.input, &ctx, extension_codec)?;
+ let x: Vec<Expr> = projection
+ .expr
+ .iter()
+ .map(|expr| expr.try_into())
+ .collect::<Result<Vec<_>, _>>()?;
+ LogicalPlanBuilder::from(input)
+ .project_with_alias(
+ x,
+ projection.optional_alias.as_ref().map(|a| match a {
+ protobuf::projection_node::OptionalAlias::Alias(alias) => {
+ alias.clone()
+ }
+ }),
+ )?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::Selection(selection) => {
+ let input: LogicalPlan =
+ into_logical_plan!(selection.input, &ctx, extension_codec)?;
+ let expr: Expr = selection
+ .expr
+ .as_ref()
+ .ok_or_else(|| {
+ BallistaError::General("expression required".to_string())
+ })?
+ .try_into()?;
+ LogicalPlanBuilder::from(input)
+ .filter(expr)?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::Window(window) => {
+ let input: LogicalPlan =
+ into_logical_plan!(window.input, &ctx, extension_codec)?;
+ let window_expr = window
+ .window_expr
+ .iter()
+ .map(|expr| expr.try_into())
+ .collect::<Result<Vec<Expr>, _>>()?;
+ LogicalPlanBuilder::from(input)
+ .window(window_expr)?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::Aggregate(aggregate) => {
+ let input: LogicalPlan =
+ into_logical_plan!(aggregate.input, &ctx, extension_codec)?;
+ let group_expr = aggregate
+ .group_expr
+ .iter()
+ .map(|expr| expr.try_into())
+ .collect::<Result<Vec<Expr>, _>>()?;
+ let aggr_expr = aggregate
+ .aggr_expr
+ .iter()
+ .map(|expr| expr.try_into())
+ .collect::<Result<Vec<Expr>, _>>()?;
+ LogicalPlanBuilder::from(input)
+ .aggregate(group_expr, aggr_expr)?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::ListingScan(scan) => {
+ let schema: Schema = convert_required!(scan.schema)?;
+
+ let mut projection = None;
+ if let Some(columns) = &scan.projection {
+ let column_indices = columns
+ .columns
+ .iter()
+ .map(|name| schema.index_of(name))
+ .collect::<Result<Vec<usize>, _>>()?;
+ projection = Some(column_indices);
+ }
+
+ let filters = scan
+ .filters
+ .iter()
+ .map(|e| e.try_into())
+ .collect::<Result<Vec<_>, _>>()?;
+
+ let file_format: Arc<dyn FileFormat> =
+ match scan.file_format_type.as_ref().ok_or_else(|| {
+ proto_error(format!(
+ "logical_plan::from_proto() Unsupported file format '{:?}'",
+ self
+ ))
+ })? {
+ &FileFormatType::Parquet(protobuf::ParquetFormat {
+ enable_pruning,
+ }) => Arc::new(
+ ParquetFormat::default().with_enable_pruning(enable_pruning),
+ ),
+ FileFormatType::Csv(protobuf::CsvFormat {
+ has_header,
+ delimiter,
+ }) => Arc::new(
+ CsvFormat::default()
+ .with_has_header(*has_header)
+ .with_delimiter(str_to_byte(delimiter)?),
+ ),
+ FileFormatType::Avro(..) => Arc::new(AvroFormat::default()),
+ };
+
+ let options = ListingOptions {
+ file_extension: scan.file_extension.clone(),
+ format: file_format,
+ table_partition_cols: scan.table_partition_cols.clone(),
+ collect_stat: scan.collect_stat,
+ target_partitions: scan.target_partitions as usize,
+ };
+
+ let object_store = ctx
+ .object_store(scan.path.as_str())
+ .map_err(|e| {
+ BallistaError::NotImplemented(format!(
+ "No object store is registered for path {}: {:?}",
+ scan.path, e
+ ))
+ })?
+ .0;
+
+ println!(
+ "Found object store {:?} for path {}",
+ object_store,
+ scan.path.as_str()
+ );
+
+ let provider = ListingTable::new(
+ object_store,
+ scan.path.clone(),
+ Arc::new(schema),
+ options,
+ );
+
+ LogicalPlanBuilder::scan_with_filters(
+ &scan.table_name,
+ Arc::new(provider),
+ projection,
+ filters,
+ )?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::Sort(sort) => {
+ let input: LogicalPlan =
+ into_logical_plan!(sort.input, &ctx, extension_codec)?;
+ let sort_expr: Vec<Expr> = sort
+ .expr
+ .iter()
+ .map(|expr| expr.try_into())
+ .collect::<Result<Vec<Expr>, _>>()?;
+ LogicalPlanBuilder::from(input)
+ .sort(sort_expr)?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::Repartition(repartition) => {
+ use datafusion::logical_plan::Partitioning;
+ let input: LogicalPlan =
+ into_logical_plan!(repartition.input, &ctx, extension_codec)?;
+ use protobuf::repartition_node::PartitionMethod;
+ let pb_partition_method = repartition.partition_method.clone().ok_or_else(|| {
+ BallistaError::General(String::from(
+ "Protobuf deserialization error, RepartitionNode was missing required field 'partition_method'",
+ ))
+ })?;
+
+ let partitioning_scheme = match pb_partition_method {
+ PartitionMethod::Hash(protobuf::HashRepartition {
+ hash_expr: pb_hash_expr,
+ partition_count,
+ }) => Partitioning::Hash(
+ pb_hash_expr
+ .iter()
+ .map(|pb_expr| pb_expr.try_into())
+ .collect::<Result<Vec<_>, _>>()?,
+ partition_count as usize,
+ ),
+ PartitionMethod::RoundRobin(partition_count) => {
+ Partitioning::RoundRobinBatch(partition_count as usize)
+ }
+ };
+
+ LogicalPlanBuilder::from(input)
+ .repartition(partitioning_scheme)?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::EmptyRelation(empty_relation) => {
+ LogicalPlanBuilder::empty(empty_relation.produce_one_row)
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::CreateExternalTable(create_extern_table) => {
+ let pb_schema = (create_extern_table.schema.clone()).ok_or_else(|| {
+ BallistaError::General(String::from(
+ "Protobuf deserialization error, CreateExternalTableNode was missing required field schema.",
+ ))
+ })?;
+
+ let pb_file_type: protobuf::FileType =
+ create_extern_table.file_type.try_into()?;
+
+ Ok(LogicalPlan::CreateExternalTable(CreateExternalTable {
+ schema: pb_schema.try_into()?,
+ name: create_extern_table.name.clone(),
+ location: create_extern_table.location.clone(),
+ file_type: pb_file_type.into(),
+ has_header: create_extern_table.has_header,
+ }))
+ }
+ LogicalPlanType::Analyze(analyze) => {
+ let input: LogicalPlan =
+ into_logical_plan!(analyze.input, &ctx, extension_codec)?;
+ LogicalPlanBuilder::from(input)
+ .explain(analyze.verbose, true)?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::Explain(explain) => {
+ let input: LogicalPlan =
+ into_logical_plan!(explain.input, &ctx, extension_codec)?;
+ LogicalPlanBuilder::from(input)
+ .explain(explain.verbose, false)?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::Limit(limit) => {
+ let input: LogicalPlan =
+ into_logical_plan!(limit.input, &ctx, extension_codec)?;
+ LogicalPlanBuilder::from(input)
+ .limit(limit.limit as usize)?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::Join(join) => {
+ let left_keys: Vec<Column> =
+ join.left_join_column.iter().map(|i| i.into()).collect();
+ let right_keys: Vec<Column> =
+ join.right_join_column.iter().map(|i| i.into()).collect();
+ let join_type =
+ protobuf::JoinType::from_i32(join.join_type).ok_or_else(|| {
+ proto_error(format!(
+ "Received a JoinNode message with unknown JoinType {}",
+ join.join_type
+ ))
+ })?;
+ let join_constraint = protobuf::JoinConstraint::from_i32(
+ join.join_constraint,
+ )
+ .ok_or_else(|| {
+ proto_error(format!(
+ "Received a JoinNode message with unknown JoinConstraint {}",
+ join.join_constraint
+ ))
+ })?;
+
+ let builder = LogicalPlanBuilder::from(into_logical_plan!(
+ join.left,
+ &ctx,
+ extension_codec
+ )?);
+ let builder = match join_constraint.into() {
+ JoinConstraint::On => builder.join(
+ &into_logical_plan!(join.right, &ctx, extension_codec)?,
+ join_type.into(),
+ (left_keys, right_keys),
+ )?,
+ JoinConstraint::Using => builder.join_using(
+ &into_logical_plan!(join.right, &ctx, extension_codec)?,
+ join_type.into(),
+ left_keys,
+ )?,
+ };
+
+ builder.build().map_err(|e| e.into())
+ }
+ LogicalPlanType::CrossJoin(crossjoin) => {
+ let left = into_logical_plan!(crossjoin.left, &ctx, extension_codec)?;
+ let right = into_logical_plan!(crossjoin.right, &ctx, extension_codec)?;
+
+ LogicalPlanBuilder::from(left)
+ .cross_join(&right)?
+ .build()
+ .map_err(|e| e.into())
+ }
+ LogicalPlanType::Extension(LogicalExtensionNode { node, inputs }) => {
+ let input_plans: Vec<LogicalPlan> = inputs
+ .iter()
+ .map(|i| i.try_into_logical_plan(ctx, extension_codec))
+ .collect::<Result<_, BallistaError>>()?;
+
+ let extension_node = extension_codec.try_decode(node, &input_plans)?;
+ Ok(LogicalPlan::Extension(extension_node))
+ }
+ }
+ }
+
+ fn try_from_logical_plan(
+ plan: &LogicalPlan,
+ extension_codec: &dyn LogicalExtensionCodec,
+ ) -> Result<Self, BallistaError>
+ where
+ Self: Sized,
+ {
+ match plan {
+ LogicalPlan::Values(Values { values, .. }) => {
+ let n_cols = if values.is_empty() {
+ 0
+ } else {
+ values[0].len()
+ } as u64;
+ let values_list = values
+ .iter()
+ .flatten()
+ .map(|v| v.try_into())
+ .collect::<Result<Vec<_>, _>>()?;
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Values(
+ protobuf::ValuesNode {
+ n_cols,
+ values_list,
+ },
+ )),
+ })
+ }
+ LogicalPlan::TableScan(TableScan {
+ table_name,
+ source,
+ filters,
+ projection,
+ ..
+ }) => {
+ let schema = source.schema();
+ let source = source.as_any();
+
+ let projection = match projection {
+ None => None,
+ Some(columns) => {
+ let column_names = columns
+ .iter()
+ .map(|i| schema.field(*i).name().to_owned())
+ .collect();
+ Some(protobuf::ProjectionColumns {
+ columns: column_names,
+ })
+ }
+ };
+ let schema: protobuf::Schema = schema.as_ref().into();
+
+ let filters: Vec<protobuf::LogicalExprNode> = filters
+ .iter()
+ .map(|filter| filter.try_into())
+ .collect::<Result<Vec<_>, _>>()?;
+
+ if let Some(listing_table) = source.downcast_ref::<ListingTable>() {
+ let any = listing_table.options().format.as_any();
+ let file_format_type = if let Some(parquet) =
+ any.downcast_ref::<ParquetFormat>()
+ {
+ FileFormatType::Parquet(protobuf::ParquetFormat {
+ enable_pruning: parquet.enable_pruning(),
+ })
+ } else if let Some(csv) = any.downcast_ref::<CsvFormat>() {
+ FileFormatType::Csv(protobuf::CsvFormat {
+ delimiter: byte_to_string(csv.delimiter())?,
+ has_header: csv.has_header(),
+ })
+ } else if any.is::<AvroFormat>() {
+ FileFormatType::Avro(protobuf::AvroFormat {})
+ } else {
+ return Err(proto_error(format!(
+ "Error converting file format, {:?} is invalid as a datafusion foramt.",
+ listing_table.options().format
+ )));
+ };
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::ListingScan(
+ protobuf::ListingTableScanNode {
+ file_format_type: Some(file_format_type),
+ table_name: table_name.to_owned(),
+ collect_stat: listing_table.options().collect_stat,
+ file_extension: listing_table
+ .options()
+ .file_extension
+ .clone(),
+ table_partition_cols: listing_table
+ .options()
+ .table_partition_cols
+ .clone(),
+ path: listing_table.table_path().to_owned(),
+ schema: Some(schema),
+ projection,
+ filters,
+ target_partitions: listing_table
+ .options()
+ .target_partitions
+ as u32,
+ },
+ )),
+ })
+ } else {
+ Err(BallistaError::General(format!(
+ "logical plan to_proto unsupported table provider {:?}",
+ source
+ )))
+ }
+ }
+ LogicalPlan::Projection(Projection {
+ expr, input, alias, ..
+ }) => Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Projection(Box::new(
+ protobuf::ProjectionNode {
+ input: Some(Box::new(
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ input.as_ref(),
+ extension_codec,
+ )?,
+ )),
+ expr: expr.iter().map(|expr| expr.try_into()).collect::<Result<
+ Vec<_>,
+ BallistaError,
+ >>(
+ )?,
+ optional_alias: alias
+ .clone()
+ .map(protobuf::projection_node::OptionalAlias::Alias),
+ },
+ ))),
+ }),
+ LogicalPlan::Filter(Filter { predicate, input }) => {
+ let input: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ input.as_ref(),
+ extension_codec,
+ )?;
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Selection(Box::new(
+ protobuf::SelectionNode {
+ input: Some(Box::new(input)),
+ expr: Some(predicate.try_into()?),
+ },
+ ))),
+ })
+ }
+ LogicalPlan::Window(Window {
+ input, window_expr, ..
+ }) => {
+ let input: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ input.as_ref(),
+ extension_codec,
+ )?;
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Window(Box::new(
+ protobuf::WindowNode {
+ input: Some(Box::new(input)),
+ window_expr: window_expr
+ .iter()
+ .map(|expr| expr.try_into())
+ .collect::<Result<Vec<_>, _>>()?,
+ },
+ ))),
+ })
+ }
+ LogicalPlan::Aggregate(Aggregate {
+ group_expr,
+ aggr_expr,
+ input,
+ ..
+ }) => {
+ let input: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ input.as_ref(),
+ extension_codec,
+ )?;
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Aggregate(Box::new(
+ protobuf::AggregateNode {
+ input: Some(Box::new(input)),
+ group_expr: group_expr
+ .iter()
+ .map(|expr| expr.try_into())
+ .collect::<Result<Vec<_>, _>>()?,
+ aggr_expr: aggr_expr
+ .iter()
+ .map(|expr| expr.try_into())
+ .collect::<Result<Vec<_>, _>>()?,
+ },
+ ))),
+ })
+ }
+ LogicalPlan::Join(Join {
+ left,
+ right,
+ on,
+ join_type,
+ join_constraint,
+ null_equals_null,
+ ..
+ }) => {
+ let left: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ left.as_ref(),
+ extension_codec,
+ )?;
+ let right: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ right.as_ref(),
+ extension_codec,
+ )?;
+ let (left_join_column, right_join_column) =
+ on.iter().map(|(l, r)| (l.into(), r.into())).unzip();
+ let join_type: protobuf::JoinType = join_type.to_owned().into();
+ let join_constraint: protobuf::JoinConstraint =
+ join_constraint.to_owned().into();
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Join(Box::new(
+ protobuf::JoinNode {
+ left: Some(Box::new(left)),
+ right: Some(Box::new(right)),
+ join_type: join_type.into(),
+ join_constraint: join_constraint.into(),
+ left_join_column,
+ right_join_column,
+ null_equals_null: *null_equals_null,
+ },
+ ))),
+ })
+ }
+ LogicalPlan::Limit(Limit { input, n }) => {
+ let input: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ input.as_ref(),
+ extension_codec,
+ )?;
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Limit(Box::new(
+ protobuf::LimitNode {
+ input: Some(Box::new(input)),
+ limit: *n as u32,
+ },
+ ))),
+ })
+ }
+ LogicalPlan::Sort(Sort { input, expr }) => {
+ let input: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ input.as_ref(),
+ extension_codec,
+ )?;
+ let selection_expr: Vec<protobuf::LogicalExprNode> = expr
+ .iter()
+ .map(|expr| expr.try_into())
+ .collect::<Result<Vec<_>, BallistaError>>()?;
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Sort(Box::new(
+ protobuf::SortNode {
+ input: Some(Box::new(input)),
+ expr: selection_expr,
+ },
+ ))),
+ })
+ }
+ LogicalPlan::Repartition(Repartition {
+ input,
+ partitioning_scheme,
+ }) => {
+ use datafusion::logical_plan::Partitioning;
+ let input: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ input.as_ref(),
+ extension_codec,
+ )?;
+
+ //Assumed common usize field was batch size
+ //Used u64 to avoid any nastyness involving large values, most data clusters are probably uniformly 64 bits any ways
+ use protobuf::repartition_node::PartitionMethod;
+
+ let pb_partition_method = match partitioning_scheme {
+ Partitioning::Hash(exprs, partition_count) => {
+ PartitionMethod::Hash(protobuf::HashRepartition {
+ hash_expr: exprs
+ .iter()
+ .map(|expr| expr.try_into())
+ .collect::<Result<Vec<_>, BallistaError>>()?,
+ partition_count: *partition_count as u64,
+ })
+ }
+ Partitioning::RoundRobinBatch(partition_count) => {
+ PartitionMethod::RoundRobin(*partition_count as u64)
+ }
+ };
+
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Repartition(Box::new(
+ protobuf::RepartitionNode {
+ input: Some(Box::new(input)),
+ partition_method: Some(pb_partition_method),
+ },
+ ))),
+ })
+ }
+ LogicalPlan::EmptyRelation(EmptyRelation {
+ produce_one_row, ..
+ }) => Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::EmptyRelation(
+ protobuf::EmptyRelationNode {
+ produce_one_row: *produce_one_row,
+ },
+ )),
+ }),
+ LogicalPlan::CreateExternalTable(CreateExternalTable {
+ name,
+ location,
+ file_type,
+ has_header,
+ schema: df_schema,
+ }) => {
+ use datafusion::sql::parser::FileType;
+
+ let pb_file_type: protobuf::FileType = match file_type {
+ FileType::NdJson => protobuf::FileType::NdJson,
+ FileType::Parquet => protobuf::FileType::Parquet,
+ FileType::CSV => protobuf::FileType::Csv,
+ FileType::Avro => protobuf::FileType::Avro,
+ };
+
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::CreateExternalTable(
+ protobuf::CreateExternalTableNode {
+ name: name.clone(),
+ location: location.clone(),
+ file_type: pb_file_type as i32,
+ has_header: *has_header,
+ schema: Some(df_schema.into()),
+ },
+ )),
+ })
+ }
+ LogicalPlan::Analyze(a) => {
+ let input = protobuf::LogicalPlanNode::try_from_logical_plan(
+ a.input.as_ref(),
+ extension_codec,
+ )?;
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Analyze(Box::new(
+ protobuf::AnalyzeNode {
+ input: Some(Box::new(input)),
+ verbose: a.verbose,
+ },
+ ))),
+ })
+ }
+ LogicalPlan::Explain(a) => {
+ let input = protobuf::LogicalPlanNode::try_from_logical_plan(
+ a.plan.as_ref(),
+ extension_codec,
+ )?;
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Explain(Box::new(
+ protobuf::ExplainNode {
+ input: Some(Box::new(input)),
+ verbose: a.verbose,
+ },
+ ))),
+ })
+ }
+ LogicalPlan::Union(_) => unimplemented!(),
+ LogicalPlan::CrossJoin(CrossJoin { left, right, .. }) => {
+ let left = protobuf::LogicalPlanNode::try_from_logical_plan(
+ left.as_ref(),
+ extension_codec,
+ )?;
+ let right = protobuf::LogicalPlanNode::try_from_logical_plan(
+ right.as_ref(),
+ extension_codec,
+ )?;
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::CrossJoin(Box::new(
+ protobuf::CrossJoinNode {
+ left: Some(Box::new(left)),
+ right: Some(Box::new(right)),
+ },
+ ))),
+ })
+ }
+ LogicalPlan::Extension(extension) => {
+ let mut buf: Vec<u8> = vec![];
+ extension_codec.try_encode(extension, &mut buf)?;
+
+ let inputs: Vec<LogicalPlanNode> = extension
+ .node
+ .inputs()
+ .iter()
+ .map(|i| {
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ i,
+ extension_codec,
+ )
+ })
+ .collect::<Result<_, BallistaError>>()?;
+
+ Ok(protobuf::LogicalPlanNode {
+ logical_plan_type: Some(LogicalPlanType::Extension(
+ LogicalExtensionNode { node: buf, inputs },
+ )),
+ })
+ }
+ LogicalPlan::CreateMemoryTable(_) => Err(proto_error(
+ "Error converting CreateMemoryTable. Not yet supported in Ballista",
+ )),
+ LogicalPlan::DropTable(_) => Err(proto_error(
+ "Error converting DropTable. Not yet supported in Ballista",
+ )),
+ }
+ }
+}
+
+#[macro_export]
+macro_rules! into_logical_plan {
+ ($PB:expr, $CTX:expr, $CODEC:expr) => {{
+ if let Some(field) = $PB.as_ref() {
+ field.as_ref().try_into_logical_plan(&$CTX, $CODEC)
+ } else {
+ Err(proto_error("Missing required field in protobuf"))
+ }
+ }};
+}
+
#[cfg(test)]
mod roundtrip_tests {
use super::super::{super::error::Result, protobuf};
use crate::error::BallistaError;
+ use crate::serde::{AsLogicalPlan, BallistaCodec};
+ use async_trait::async_trait;
use core::panic;
+ use datafusion::datasource::listing::ListingTable;
+ use datafusion::datasource::object_store::{
+ FileMetaStream, ListEntryStream, ObjectReader, ObjectStore, SizedFile,
+ };
+ use datafusion::datasource::TableProvider;
+ use datafusion::error::DataFusionError;
use datafusion::{
arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit, UnionMode},
datasource::object_store::local::LocalFileSystem,
@@ -37,8 +871,43 @@ mod roundtrip_tests {
sql::parser::FileType,
};
use protobuf::arrow_type;
+ use sqlparser::test_utils::table;
use std::{convert::TryInto, sync::Arc};
+ #[derive(Debug)]
+ struct TestObjectStore {}
+
+ #[async_trait]
+ impl ObjectStore for TestObjectStore {
+ async fn list_file(
+ &self,
+ _prefix: &str,
+ ) -> datafusion::error::Result<FileMetaStream> {
+ Err(DataFusionError::NotImplemented(
+ "this is only a test object store".to_string(),
+ ))
+ }
+
+ async fn list_dir(
+ &self,
+ _prefix: &str,
+ _delimiter: Option<String>,
+ ) -> datafusion::error::Result<ListEntryStream> {
+ Err(DataFusionError::NotImplemented(
+ "this is only a test object store".to_string(),
+ ))
+ }
+
+ fn file_reader(
+ &self,
+ _file: SizedFile,
+ ) -> datafusion::error::Result<Arc<dyn ObjectReader>> {
+ Err(DataFusionError::NotImplemented(
+ "this is only a test object store".to_string(),
+ ))
+ }
+ }
+
//Given a identity of a LogicalPlan converts it to protobuf and back, using debug formatting to test equality.
macro_rules! roundtrip_test {
($initial_struct:ident, $proto_type:ty, $struct_type:ty) => {
@@ -55,7 +924,42 @@ mod roundtrip_tests {
roundtrip_test!($initial_struct, protobuf::LogicalPlanNode, $struct_type);
};
($initial_struct:ident) => {
- roundtrip_test!($initial_struct, protobuf::LogicalPlanNode, LogicalPlan);
+ let ctx = ExecutionContext::new();
+ let codec: BallistaCodec<
+ protobuf::LogicalPlanNode,
+ protobuf::PhysicalPlanNode,
+ > = BallistaCodec::default();
+ let proto: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ &$initial_struct,
+ codec.logical_extension_codec(),
+ )
+ .expect("from logical plan");
+ let round_trip: LogicalPlan = proto
+ .try_into_logical_plan(&ctx, codec.logical_extension_codec())
+ .expect("to logical plan");
+
+ assert_eq!(
+ format!("{:?}", $initial_struct),
+ format!("{:?}", round_trip)
+ );
+ };
+ ($initial_struct:ident, $ctx:ident) => {
+ let codec: BallistaCodec<
+ protobuf::LogicalPlanNode,
+ protobuf::PhysicalPlanNode,
+ > = BallistaCodec::default();
+ let proto: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(&$initial_struct)
+ .expect("from logical plan");
+ let round_trip: LogicalPlan = proto
+ .try_into_logical_plan(&$ctx, codec.logical_extension_codec())
+ .expect("to logical plan");
+
+ assert_eq!(
+ format!("{:?}", $initial_struct),
+ format!("{:?}", round_trip)
+ );
};
}
@@ -876,6 +1780,66 @@ mod roundtrip_tests {
Ok(())
}
+ #[tokio::test]
+ async fn roundtrip_logical_plan_custom_ctx() -> Result<()> {
+ let ctx = ExecutionContext::new();
+ let codec: BallistaCodec<protobuf::LogicalPlanNode, protobuf::PhysicalPlanNode> =
+ BallistaCodec::default();
+ let custom_object_store = Arc::new(TestObjectStore {});
+ ctx.register_object_store("test", custom_object_store.clone());
+
+ let (os, _) = ctx.object_store("test://foo.csv")?;
+
+ println!("Object Store {:?}", os);
+
+ let schema = Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("first_name", DataType::Utf8, false),
+ Field::new("last_name", DataType::Utf8, false),
+ Field::new("state", DataType::Utf8, false),
+ Field::new("salary", DataType::Int32, false),
+ ]);
+
+ let plan = LogicalPlanBuilder::scan_csv(
+ custom_object_store.clone(),
+ "test://employee.csv",
+ CsvReadOptions::new().schema(&schema).has_header(true),
+ Some(vec![3, 4]),
+ 4,
+ )
+ .await
+ .and_then(|plan| plan.build())
+ .map_err(BallistaError::DataFusionError)?;
+
+ let proto: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ &plan,
+ codec.logical_extension_codec(),
+ )
+ .expect("from logical plan");
+ let round_trip: LogicalPlan = proto
+ .try_into_logical_plan(&ctx, codec.logical_extension_codec())
+ .expect("to logical plan");
+
+ assert_eq!(format!("{:?}", plan), format!("{:?}", round_trip));
+
+ let round_trip_store = match round_trip {
+ LogicalPlan::TableScan(scan) => {
+ match scan.source.as_ref().as_any().downcast_ref::<ListingTable>() {
+ Some(listing_table) => {
+ format!("{:?}", listing_table.object_store())
+ }
+ _ => panic!("expected a ListingTable"),
+ }
+ }
+ _ => panic!("expected a TableScan"),
+ };
+
+ assert_eq!(round_trip_store, format!("{:?}", custom_object_store));
+
+ Ok(())
+ }
+
#[test]
fn roundtrip_not() -> Result<()> {
let test_expr = Expr::Not(Box::new(Expr::Literal((1.0).into())));
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 84910b2..16f19b0 100644
--- a/ballista/rust/core/src/serde/logical_plan/to_proto.rs
+++ b/ballista/rust/core/src/serde/logical_plan/to_proto.rs
@@ -619,355 +619,6 @@ impl TryFrom<&datafusion::scalar::ScalarValue> for protobuf::ScalarValue {
}
}
-impl TryInto<protobuf::LogicalPlanNode> for &LogicalPlan {
- type Error = BallistaError;
-
- fn try_into(self) -> Result<protobuf::LogicalPlanNode, Self::Error> {
- use protobuf::logical_plan_node::LogicalPlanType;
- match self {
- LogicalPlan::Values(Values { values, .. }) => {
- let n_cols = if values.is_empty() {
- 0
- } else {
- values[0].len()
- } as u64;
- let values_list = values
- .iter()
- .flatten()
- .map(|v| v.try_into())
- .collect::<Result<Vec<_>, _>>()?;
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Values(
- protobuf::ValuesNode {
- n_cols,
- values_list,
- },
- )),
- })
- }
- LogicalPlan::TableScan(TableScan {
- table_name,
- source,
- filters,
- projection,
- ..
- }) => {
- let schema = source.schema();
- let source = source.as_any();
-
- let projection = match projection {
- None => None,
- Some(columns) => {
- let column_names = columns
- .iter()
- .map(|i| schema.field(*i).name().to_owned())
- .collect();
- Some(protobuf::ProjectionColumns {
- columns: column_names,
- })
- }
- };
- let schema: protobuf::Schema = schema.as_ref().into();
-
- let filters: Vec<protobuf::LogicalExprNode> = filters
- .iter()
- .map(|filter| filter.try_into())
- .collect::<Result<Vec<_>, _>>()?;
-
- if let Some(listing_table) = source.downcast_ref::<ListingTable>() {
- let any = listing_table.options().format.as_any();
- let file_format_type = if let Some(parquet) =
- any.downcast_ref::<ParquetFormat>()
- {
- FileFormatType::Parquet(protobuf::ParquetFormat {
- enable_pruning: parquet.enable_pruning(),
- })
- } else if let Some(csv) = any.downcast_ref::<CsvFormat>() {
- FileFormatType::Csv(protobuf::CsvFormat {
- delimiter: byte_to_string(csv.delimiter())?,
- has_header: csv.has_header(),
- })
- } else if any.is::<AvroFormat>() {
- FileFormatType::Avro(protobuf::AvroFormat {})
- } else {
- return Err(proto_error(format!(
- "Error converting file format, {:?} is invalid as a datafusion foramt.",
- listing_table.options().format
- )));
- };
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::ListingScan(
- protobuf::ListingTableScanNode {
- file_format_type: Some(file_format_type),
- table_name: table_name.to_owned(),
- collect_stat: listing_table.options().collect_stat,
- file_extension: listing_table
- .options()
- .file_extension
- .clone(),
- table_partition_cols: listing_table
- .options()
- .table_partition_cols
- .clone(),
- path: listing_table.table_path().to_owned(),
- schema: Some(schema),
- projection,
- filters,
- target_partitions: listing_table
- .options()
- .target_partitions
- as u32,
- },
- )),
- })
- } else {
- Err(BallistaError::General(format!(
- "logical plan to_proto unsupported table provider {:?}",
- source
- )))
- }
- }
- LogicalPlan::Projection(Projection {
- expr, input, alias, ..
- }) => Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Projection(Box::new(
- protobuf::ProjectionNode {
- input: Some(Box::new(input.as_ref().try_into()?)),
- expr: expr.iter().map(|expr| expr.try_into()).collect::<Result<
- Vec<_>,
- BallistaError,
- >>(
- )?,
- optional_alias: alias
- .clone()
- .map(protobuf::projection_node::OptionalAlias::Alias),
- },
- ))),
- }),
- LogicalPlan::Filter(Filter { predicate, input }) => {
- let input: protobuf::LogicalPlanNode = input.as_ref().try_into()?;
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Selection(Box::new(
- protobuf::SelectionNode {
- input: Some(Box::new(input)),
- expr: Some(predicate.try_into()?),
- },
- ))),
- })
- }
- LogicalPlan::Window(Window {
- input, window_expr, ..
- }) => {
- let input: protobuf::LogicalPlanNode = input.as_ref().try_into()?;
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Window(Box::new(
- protobuf::WindowNode {
- input: Some(Box::new(input)),
- window_expr: window_expr
- .iter()
- .map(|expr| expr.try_into())
- .collect::<Result<Vec<_>, _>>()?,
- },
- ))),
- })
- }
- LogicalPlan::Aggregate(Aggregate {
- group_expr,
- aggr_expr,
- input,
- ..
- }) => {
- let input: protobuf::LogicalPlanNode = input.as_ref().try_into()?;
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Aggregate(Box::new(
- protobuf::AggregateNode {
- input: Some(Box::new(input)),
- group_expr: group_expr
- .iter()
- .map(|expr| expr.try_into())
- .collect::<Result<Vec<_>, _>>()?,
- aggr_expr: aggr_expr
- .iter()
- .map(|expr| expr.try_into())
- .collect::<Result<Vec<_>, _>>()?,
- },
- ))),
- })
- }
- LogicalPlan::Join(Join {
- left,
- right,
- on,
- join_type,
- join_constraint,
- null_equals_null,
- ..
- }) => {
- let left: protobuf::LogicalPlanNode = left.as_ref().try_into()?;
- let right: protobuf::LogicalPlanNode = right.as_ref().try_into()?;
- let (left_join_column, right_join_column) =
- on.iter().map(|(l, r)| (l.into(), r.into())).unzip();
- let join_type: protobuf::JoinType = join_type.to_owned().into();
- let join_constraint: protobuf::JoinConstraint =
- join_constraint.to_owned().into();
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Join(Box::new(
- protobuf::JoinNode {
- left: Some(Box::new(left)),
- right: Some(Box::new(right)),
- join_type: join_type.into(),
- join_constraint: join_constraint.into(),
- left_join_column,
- right_join_column,
- null_equals_null: *null_equals_null,
- },
- ))),
- })
- }
- LogicalPlan::Limit(Limit { input, n }) => {
- let input: protobuf::LogicalPlanNode = input.as_ref().try_into()?;
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Limit(Box::new(
- protobuf::LimitNode {
- input: Some(Box::new(input)),
- limit: *n as u32,
- },
- ))),
- })
- }
- LogicalPlan::Sort(Sort { input, expr }) => {
- let input: protobuf::LogicalPlanNode = input.as_ref().try_into()?;
- let selection_expr: Vec<protobuf::LogicalExprNode> = expr
- .iter()
- .map(|expr| expr.try_into())
- .collect::<Result<Vec<_>, BallistaError>>()?;
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Sort(Box::new(
- protobuf::SortNode {
- input: Some(Box::new(input)),
- expr: selection_expr,
- },
- ))),
- })
- }
- LogicalPlan::Repartition(Repartition {
- input,
- partitioning_scheme,
- }) => {
- use datafusion::logical_plan::Partitioning;
- let input: protobuf::LogicalPlanNode = input.as_ref().try_into()?;
-
- //Assumed common usize field was batch size
- //Used u64 to avoid any nastyness involving large values, most data clusters are probably uniformly 64 bits any ways
- use protobuf::repartition_node::PartitionMethod;
-
- let pb_partition_method = match partitioning_scheme {
- Partitioning::Hash(exprs, partition_count) => {
- PartitionMethod::Hash(protobuf::HashRepartition {
- hash_expr: exprs
- .iter()
- .map(|expr| expr.try_into())
- .collect::<Result<Vec<_>, BallistaError>>()?,
- partition_count: *partition_count as u64,
- })
- }
- Partitioning::RoundRobinBatch(partition_count) => {
- PartitionMethod::RoundRobin(*partition_count as u64)
- }
- };
-
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Repartition(Box::new(
- protobuf::RepartitionNode {
- input: Some(Box::new(input)),
- partition_method: Some(pb_partition_method),
- },
- ))),
- })
- }
- LogicalPlan::EmptyRelation(EmptyRelation {
- produce_one_row, ..
- }) => Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::EmptyRelation(
- protobuf::EmptyRelationNode {
- produce_one_row: *produce_one_row,
- },
- )),
- }),
- LogicalPlan::CreateExternalTable(CreateExternalTable {
- name,
- location,
- file_type,
- has_header,
- schema: df_schema,
- }) => {
- use datafusion::sql::parser::FileType;
-
- let pb_file_type: protobuf::FileType = match file_type {
- FileType::NdJson => protobuf::FileType::NdJson,
- FileType::Parquet => protobuf::FileType::Parquet,
- FileType::CSV => protobuf::FileType::Csv,
- FileType::Avro => protobuf::FileType::Avro,
- };
-
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::CreateExternalTable(
- protobuf::CreateExternalTableNode {
- name: name.clone(),
- location: location.clone(),
- file_type: pb_file_type as i32,
- has_header: *has_header,
- schema: Some(df_schema.into()),
- },
- )),
- })
- }
- LogicalPlan::Analyze(a) => {
- let input: protobuf::LogicalPlanNode = a.input.as_ref().try_into()?;
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Analyze(Box::new(
- protobuf::AnalyzeNode {
- input: Some(Box::new(input)),
- verbose: a.verbose,
- },
- ))),
- })
- }
- LogicalPlan::Explain(a) => {
- let input: protobuf::LogicalPlanNode = a.plan.as_ref().try_into()?;
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::Explain(Box::new(
- protobuf::ExplainNode {
- input: Some(Box::new(input)),
- verbose: a.verbose,
- },
- ))),
- })
- }
- LogicalPlan::Extension { .. } => unimplemented!(),
- LogicalPlan::Union(_) => unimplemented!(),
- LogicalPlan::CrossJoin(CrossJoin { left, right, .. }) => {
- let left: protobuf::LogicalPlanNode = left.as_ref().try_into()?;
- let right: protobuf::LogicalPlanNode = right.as_ref().try_into()?;
- Ok(protobuf::LogicalPlanNode {
- logical_plan_type: Some(LogicalPlanType::CrossJoin(Box::new(
- protobuf::CrossJoinNode {
- left: Some(Box::new(left)),
- right: Some(Box::new(right)),
- },
- ))),
- })
- }
- LogicalPlan::CreateMemoryTable(_) => Err(proto_error(
- "Error converting CreateMemoryTable. Not yet supported in Ballista",
- )),
- LogicalPlan::DropTable(_) => Err(proto_error(
- "Error converting DropTable. Not yet supported in Ballista",
- )),
- }
- }
-}
-
fn create_proto_scalar<I, T: FnOnce(&I) -> protobuf::scalar_value::Value>(
v: &Option<I>,
null_arrow_type: protobuf::PrimitiveScalarType,
diff --git a/ballista/rust/core/src/serde/mod.rs b/ballista/rust/core/src/serde/mod.rs
index f7b0b94..3a1736c 100644
--- a/ballista/rust/core/src/serde/mod.rs
+++ b/ballista/rust/core/src/serde/mod.rs
@@ -18,15 +18,24 @@
//! This crate contains code generated from the Ballista Protocol Buffer Definition as well
//! as convenience code for interacting with the generated code.
+use prost::bytes::{Buf, BufMut};
+use std::fmt::Debug;
+use std::marker::PhantomData;
+use std::sync::Arc;
use std::{convert::TryInto, io::Cursor};
use datafusion::arrow::datatypes::{IntervalUnit, UnionMode};
-use datafusion::logical_plan::{JoinConstraint, JoinType, Operator};
+use datafusion::logical_plan::{
+ JoinConstraint, JoinType, LogicalPlan, Operator, UserDefinedLogicalNode,
+};
use datafusion::physical_plan::aggregates::AggregateFunction;
use datafusion::physical_plan::window_functions::BuiltInWindowFunction;
use crate::{error::BallistaError, serde::scheduler::Action as BallistaAction};
+use datafusion::logical_plan::plan::Extension;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::ExecutionContext;
use prost::Message;
// include the generated protobuf source as a submodule
@@ -51,6 +60,175 @@ pub(crate) fn proto_error<S: Into<String>>(message: S) -> BallistaError {
BallistaError::General(message.into())
}
+pub trait AsLogicalPlan: Debug + Send + Sync + Clone {
+ fn try_decode(buf: &[u8]) -> Result<Self, BallistaError>
+ where
+ Self: Sized;
+
+ fn try_encode<B>(&self, buf: &mut B) -> Result<(), BallistaError>
+ where
+ B: BufMut,
+ Self: Sized;
+
+ fn try_into_logical_plan(
+ &self,
+ ctx: &ExecutionContext,
+ extension_codec: &dyn LogicalExtensionCodec,
+ ) -> Result<LogicalPlan, BallistaError>;
+
+ fn try_from_logical_plan(
+ plan: &LogicalPlan,
+ extension_codec: &dyn LogicalExtensionCodec,
+ ) -> Result<Self, BallistaError>
+ where
+ Self: Sized;
+}
+
+pub trait LogicalExtensionCodec: Debug + Send + Sync {
+ fn try_decode(
+ &self,
+ buf: &[u8],
+ inputs: &[LogicalPlan],
+ ) -> Result<Extension, BallistaError>;
+
+ fn try_encode(
+ &self,
+ node: &Extension,
+ buf: &mut Vec<u8>,
+ ) -> Result<(), BallistaError>;
+}
+
+#[derive(Debug, Clone)]
+pub struct DefaultLogicalExtensionCodec {}
+
+impl LogicalExtensionCodec for DefaultLogicalExtensionCodec {
+ fn try_decode(
+ &self,
+ _buf: &[u8],
+ _inputs: &[LogicalPlan],
+ ) -> Result<Extension, BallistaError> {
+ Err(BallistaError::NotImplemented(
+ "LogicalExtensionCodec is not provided".to_string(),
+ ))
+ }
+
+ fn try_encode(
+ &self,
+ _node: &Extension,
+ _buf: &mut Vec<u8>,
+ ) -> Result<(), BallistaError> {
+ Err(BallistaError::NotImplemented(
+ "LogicalExtensionCodec is not provided".to_string(),
+ ))
+ }
+}
+
+pub trait AsExecutionPlan: Debug + Send + Sync + Clone {
+ fn try_decode(buf: &[u8]) -> Result<Self, BallistaError>
+ where
+ Self: Sized;
+
+ fn try_encode<B>(&self, buf: &mut B) -> Result<(), BallistaError>
+ where
+ B: BufMut,
+ Self: Sized;
+
+ fn try_into_physical_plan(
+ &self,
+ ctx: &ExecutionContext,
+ extension_codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<Arc<dyn ExecutionPlan>, BallistaError>;
+
+ fn try_from_physical_plan(
+ plan: Arc<dyn ExecutionPlan>,
+ extension_codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<Self, BallistaError>
+ where
+ Self: Sized;
+}
+
+pub trait PhysicalExtensionCodec: Debug + Send + Sync {
+ fn try_decode(
+ &self,
+ buf: &[u8],
+ inputs: &[Arc<dyn ExecutionPlan>],
+ ) -> Result<Arc<dyn ExecutionPlan>, BallistaError>;
+
+ fn try_encode(
+ &self,
+ node: Arc<dyn ExecutionPlan>,
+ buf: &mut Vec<u8>,
+ ) -> Result<(), BallistaError>;
+}
+
+#[derive(Debug, Clone)]
+pub struct DefaultPhysicalExtensionCodec {}
+
+impl PhysicalExtensionCodec for DefaultPhysicalExtensionCodec {
+ fn try_decode(
+ &self,
+ _buf: &[u8],
+ _inputs: &[Arc<dyn ExecutionPlan>],
+ ) -> Result<Arc<dyn ExecutionPlan>, BallistaError> {
+ Err(BallistaError::NotImplemented(
+ "PhysicalExtensionCodec is not provided".to_string(),
+ ))
+ }
+
+ fn try_encode(
+ &self,
+ _node: Arc<dyn ExecutionPlan>,
+ _buf: &mut Vec<u8>,
+ ) -> Result<(), BallistaError> {
+ Err(BallistaError::NotImplemented(
+ "PhysicalExtensionCodec is not provided".to_string(),
+ ))
+ }
+}
+
+#[derive(Clone, Debug)]
+pub struct BallistaCodec<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> {
+ logical_extension_codec: Arc<dyn LogicalExtensionCodec>,
+ physical_extension_codec: Arc<dyn PhysicalExtensionCodec>,
+ logical_plan_repr: PhantomData<T>,
+ physical_plan_repr: PhantomData<U>,
+}
+
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> Default
+ for BallistaCodec<T, U>
+{
+ fn default() -> Self {
+ Self {
+ logical_extension_codec: Arc::new(DefaultLogicalExtensionCodec {}),
+ physical_extension_codec: Arc::new(DefaultPhysicalExtensionCodec {}),
+ logical_plan_repr: PhantomData,
+ physical_plan_repr: PhantomData,
+ }
+ }
+}
+
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> BallistaCodec<T, U> {
+ pub fn new(
+ logical_extension_codec: Arc<dyn LogicalExtensionCodec>,
+ physical_extension_codec: Arc<dyn PhysicalExtensionCodec>,
+ ) -> Self {
+ Self {
+ logical_extension_codec,
+ physical_extension_codec,
+ logical_plan_repr: PhantomData,
+ physical_plan_repr: PhantomData,
+ }
+ }
+
+ pub fn logical_extension_codec(&self) -> &dyn LogicalExtensionCodec {
+ self.logical_extension_codec.as_ref()
+ }
+
+ pub fn physical_extension_codec(&self) -> &dyn PhysicalExtensionCodec {
+ self.physical_extension_codec.as_ref()
+ }
+}
+
#[macro_export]
macro_rules! convert_required {
($PB:expr) => {{
@@ -400,3 +578,402 @@ fn vec_to_array<T, const N: usize>(v: Vec<T>) -> [T; N] {
panic!("Expected a Vec of length {} but it was {}", N, v.len())
})
}
+
+#[cfg(test)]
+mod tests {
+ use async_trait::async_trait;
+ use datafusion::arrow::datatypes::SchemaRef;
+ use datafusion::datasource::object_store::local::LocalFileSystem;
+ use datafusion::error::DataFusionError;
+ use datafusion::execution::context::{ExecutionContextState, QueryPlanner};
+ use datafusion::execution::runtime_env::RuntimeEnv;
+ use datafusion::logical_plan::plan::Extension;
+ use datafusion::logical_plan::{
+ col, DFSchemaRef, Expr, LogicalPlan, LogicalPlanBuilder, UserDefinedLogicalNode,
+ };
+ use datafusion::physical_plan::planner::{DefaultPhysicalPlanner, ExtensionPlanner};
+ use datafusion::physical_plan::{
+ DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PhysicalPlanner,
+ SendableRecordBatchStream, Statistics,
+ };
+ use datafusion::prelude::{CsvReadOptions, ExecutionConfig, ExecutionContext};
+ use prost::Message;
+ use std::any::Any;
+ use std::collections::BTreeMap;
+ use std::convert::TryInto;
+ use std::fmt;
+ use std::fmt::{Debug, Formatter};
+ use std::sync::Arc;
+
+ pub mod proto {
+ use crate::serde::protobuf;
+ use prost::Message;
+
+ #[derive(Clone, PartialEq, ::prost::Message)]
+ pub struct TopKPlanProto {
+ #[prost(uint64, tag = "1")]
+ pub k: u64,
+
+ #[prost(message, optional, tag = "2")]
+ pub expr: ::core::option::Option<protobuf::LogicalExprNode>,
+ }
+
+ #[derive(Clone, PartialEq, ::prost::Message)]
+ pub struct TopKExecProto {
+ #[prost(uint64, tag = "1")]
+ pub k: u64,
+ }
+ }
+
+ use crate::error::BallistaError;
+ use crate::serde::protobuf::{LogicalPlanNode, PhysicalPlanNode};
+ use crate::serde::{
+ AsExecutionPlan, AsLogicalPlan, BallistaCodec, LogicalExtensionCodec,
+ PhysicalExtensionCodec,
+ };
+ use proto::{TopKExecProto, TopKPlanProto};
+
+ struct TopKPlanNode {
+ k: usize,
+ input: LogicalPlan,
+ /// The sort expression (this example only supports a single sort
+ /// expr)
+ expr: Expr,
+ }
+
+ impl TopKPlanNode {
+ pub fn new(k: usize, input: LogicalPlan, expr: Expr) -> Self {
+ Self { k, input, expr }
+ }
+ }
+
+ impl Debug for TopKPlanNode {
+ fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+ self.fmt_for_explain(f)
+ }
+ }
+
+ impl UserDefinedLogicalNode for TopKPlanNode {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn inputs(&self) -> Vec<&LogicalPlan> {
+ vec![&self.input]
+ }
+
+ /// Schema for TopK is the same as the input
+ fn schema(&self) -> &DFSchemaRef {
+ self.input.schema()
+ }
+
+ fn expressions(&self) -> Vec<Expr> {
+ vec![self.expr.clone()]
+ }
+
+ /// For example: `TopK: k=10`
+ fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "TopK: k={}", self.k)
+ }
+
+ fn from_template(
+ &self,
+ exprs: &[Expr],
+ inputs: &[LogicalPlan],
+ ) -> Arc<dyn UserDefinedLogicalNode + Send + Sync> {
+ assert_eq!(inputs.len(), 1, "input size inconsistent");
+ assert_eq!(exprs.len(), 1, "expression size inconsistent");
+ Arc::new(TopKPlanNode {
+ k: self.k,
+ input: inputs[0].clone(),
+ expr: exprs[0].clone(),
+ })
+ }
+ }
+
+ struct TopKExec {
+ input: Arc<dyn ExecutionPlan>,
+ /// The maxium number of values
+ k: usize,
+ }
+
+ impl TopKExec {
+ pub fn new(k: usize, input: Arc<dyn ExecutionPlan>) -> Self {
+ Self { input, k }
+ }
+ }
+
+ impl Debug for TopKExec {
+ fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+ write!(f, "TopKExec")
+ }
+ }
+
+ #[async_trait]
+ impl ExecutionPlan for TopKExec {
+ /// Return a reference to Any that can be used for downcasting
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn schema(&self) -> SchemaRef {
+ self.input.schema()
+ }
+
+ fn output_partitioning(&self) -> Partitioning {
+ Partitioning::UnknownPartitioning(1)
+ }
+
+ fn required_child_distribution(&self) -> Distribution {
+ Distribution::SinglePartition
+ }
+
+ fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
+ vec![self.input.clone()]
+ }
+
+ fn with_new_children(
+ &self,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+ match children.len() {
+ 1 => Ok(Arc::new(TopKExec {
+ input: children[0].clone(),
+ k: self.k,
+ })),
+ _ => Err(DataFusionError::Internal(
+ "TopKExec wrong number of children".to_string(),
+ )),
+ }
+ }
+
+ /// Execute one partition and return an iterator over RecordBatch
+ async fn execute(
+ &self,
+ _partition: usize,
+ _runtime: Arc<RuntimeEnv>,
+ ) -> datafusion::error::Result<SendableRecordBatchStream> {
+ Err(DataFusionError::NotImplemented(
+ "not implemented".to_string(),
+ ))
+ }
+
+ fn fmt_as(
+ &self,
+ t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default => {
+ write!(f, "TopKExec: k={}", self.k)
+ }
+ }
+ }
+
+ fn statistics(&self) -> Statistics {
+ // to improve the optimizability of this plan
+ // better statistics inference could be provided
+ Statistics::default()
+ }
+ }
+
+ struct TopKPlanner {}
+
+ impl ExtensionPlanner for TopKPlanner {
+ /// Create a physical plan for an extension node
+ fn plan_extension(
+ &self,
+ _planner: &dyn PhysicalPlanner,
+ node: &dyn UserDefinedLogicalNode,
+ logical_inputs: &[&LogicalPlan],
+ physical_inputs: &[Arc<dyn ExecutionPlan>],
+ _ctx_state: &ExecutionContextState,
+ ) -> datafusion::error::Result<Option<Arc<dyn ExecutionPlan>>> {
+ Ok(
+ if let Some(topk_node) = node.as_any().downcast_ref::<TopKPlanNode>() {
+ assert_eq!(logical_inputs.len(), 1, "Inconsistent number of inputs");
+ assert_eq!(physical_inputs.len(), 1, "Inconsistent number of inputs");
+ // figure out input name
+ Some(Arc::new(TopKExec {
+ input: physical_inputs[0].clone(),
+ k: topk_node.k,
+ }))
+ } else {
+ None
+ },
+ )
+ }
+ }
+
+ struct TopKQueryPlanner {}
+
+ #[async_trait]
+ impl QueryPlanner for TopKQueryPlanner {
+ /// Given a `LogicalPlan` created from above, create an
+ /// `ExecutionPlan` suitable for execution
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ ctx_state: &ExecutionContextState,
+ ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+ // Teach the default physical planner how to plan TopK nodes.
+ let physical_planner =
+ DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
+ TopKPlanner {},
+ )]);
+ // Delegate most work of physical planning to the default physical planner
+ physical_planner
+ .create_physical_plan(logical_plan, ctx_state)
+ .await
+ }
+ }
+
+ #[derive(Debug)]
+ pub struct TopKExtensionCodec {}
+
+ impl LogicalExtensionCodec for TopKExtensionCodec {
+ fn try_decode(
+ &self,
+ buf: &[u8],
+ inputs: &[LogicalPlan],
+ ) -> Result<Extension, BallistaError> {
+ if let Some((input, _)) = inputs.split_first() {
+ let proto = TopKPlanProto::decode(buf).map_err(|e| {
+ BallistaError::Internal(format!(
+ "failed to decode logical plan: {:?}",
+ e
+ ))
+ })?;
+
+ if let Some(expr) = proto.expr.as_ref() {
+ let node = TopKPlanNode::new(
+ proto.k as usize,
+ input.clone(),
+ expr.try_into()?,
+ );
+
+ Ok(Extension {
+ node: Arc::new(node),
+ })
+ } else {
+ Err(BallistaError::from("invalid plan, no expr".to_string()))
+ }
+ } else {
+ Err(BallistaError::from("invalid plan, no input".to_string()))
+ }
+ }
+
+ fn try_encode(
+ &self,
+ node: &Extension,
+ buf: &mut Vec<u8>,
+ ) -> Result<(), BallistaError> {
+ if let Some(exec) = node.node.as_any().downcast_ref::<TopKPlanNode>() {
+ let proto = TopKPlanProto {
+ k: exec.k as u64,
+ expr: Some((&exec.expr).try_into()?),
+ };
+
+ proto.encode(buf).map_err(|e| {
+ BallistaError::Internal(format!(
+ "failed to encode logical plan: {:?}",
+ e
+ ))
+ })?;
+
+ Ok(())
+ } else {
+ Err(BallistaError::from("unsupported plan type".to_string()))
+ }
+ }
+ }
+
+ impl PhysicalExtensionCodec for TopKExtensionCodec {
+ fn try_decode(
+ &self,
+ buf: &[u8],
+ inputs: &[Arc<dyn ExecutionPlan>],
+ ) -> Result<Arc<dyn ExecutionPlan>, BallistaError> {
+ if let Some((input, _)) = inputs.split_first() {
+ let proto = TopKExecProto::decode(buf).map_err(|e| {
+ BallistaError::Internal(format!(
+ "failed to decode execution plan: {:?}",
+ e
+ ))
+ })?;
+ Ok(Arc::new(TopKExec::new(proto.k as usize, input.clone())))
+ } else {
+ Err(BallistaError::from("invalid plan, no input".to_string()))
+ }
+ }
+
+ fn try_encode(
+ &self,
+ node: Arc<dyn ExecutionPlan>,
+ buf: &mut Vec<u8>,
+ ) -> Result<(), BallistaError> {
+ if let Some(exec) = node.as_any().downcast_ref::<TopKExec>() {
+ let proto = TopKExecProto { k: exec.k as u64 };
+
+ proto.encode(buf).map_err(|e| {
+ BallistaError::Internal(format!(
+ "failed to encode execution plan: {:?}",
+ e
+ ))
+ })?;
+
+ Ok(())
+ } else {
+ Err(BallistaError::from("unsupported plan type".to_string()))
+ }
+ }
+ }
+
+ #[tokio::test]
+ async fn test_extension_plan() -> crate::error::Result<()> {
+ let store = Arc::new(LocalFileSystem {});
+ let config =
+ ExecutionConfig::new().with_query_planner(Arc::new(TopKQueryPlanner {}));
+
+ let ctx = ExecutionContext::with_config(config);
+
+ let scan = LogicalPlanBuilder::scan_csv(
+ store,
+ "../../../datafusion/tests/customer.csv",
+ CsvReadOptions::default(),
+ None,
+ 1,
+ )
+ .await?
+ .build()?;
+
+ let topk_plan = LogicalPlan::Extension(Extension {
+ node: Arc::new(TopKPlanNode::new(3, scan, col("revenue"))),
+ });
+
+ let topk_exec = ctx.create_physical_plan(&topk_plan).await?;
+
+ let extension_codec = TopKExtensionCodec {};
+
+ let proto = LogicalPlanNode::try_from_logical_plan(&topk_plan, &extension_codec)?;
+ let logical_round_trip = proto.try_into_logical_plan(&ctx, &extension_codec)?;
+
+ assert_eq!(
+ format!("{:?}", topk_plan),
+ format!("{:?}", logical_round_trip)
+ );
+
+ let proto = PhysicalPlanNode::try_from_physical_plan(
+ topk_exec.clone(),
+ &extension_codec,
+ )?;
+ let physical_round_trip = proto.try_into_physical_plan(&ctx, &extension_codec)?;
+
+ assert_eq!(
+ format!("{:?}", topk_exec),
+ format!("{:?}", physical_round_trip)
+ );
+
+ Ok(())
+ }
+}
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 5dd57d1..93629a3 100644
--- a/ballista/rust/core/src/serde/physical_plan/from_proto.rs
+++ b/ballista/rust/core/src/serde/physical_plan/from_proto.rs
@@ -84,388 +84,6 @@ use log::debug;
use protobuf::physical_expr_node::ExprType;
use protobuf::physical_plan_node::PhysicalPlanType;
-impl TryInto<Arc<dyn ExecutionPlan>> for &protobuf::PhysicalPlanNode {
- type Error = BallistaError;
-
- fn try_into(self) -> Result<Arc<dyn ExecutionPlan>, Self::Error> {
- let plan = self.physical_plan_type.as_ref().ok_or_else(|| {
- proto_error(format!(
- "physical_plan::from_proto() Unsupported physical plan '{:?}'",
- self
- ))
- })?;
- match plan {
- PhysicalPlanType::Projection(projection) => {
- let input: Arc<dyn ExecutionPlan> =
- convert_box_required!(projection.input)?;
- let exprs = projection
- .expr
- .iter()
- .zip(projection.expr_name.iter())
- .map(|(expr, name)| Ok((expr.try_into()?, name.to_string())))
- .collect::<Result<Vec<(Arc<dyn PhysicalExpr>, String)>, Self::Error>>(
- )?;
- Ok(Arc::new(ProjectionExec::try_new(exprs, input)?))
- }
- PhysicalPlanType::Filter(filter) => {
- let input: Arc<dyn ExecutionPlan> = convert_box_required!(filter.input)?;
- let predicate = filter
- .expr
- .as_ref()
- .ok_or_else(|| {
- BallistaError::General(
- "filter (FilterExecNode) in PhysicalPlanNode is missing."
- .to_owned(),
- )
- })?
- .try_into()?;
- Ok(Arc::new(FilterExec::try_new(predicate, input)?))
- }
- PhysicalPlanType::CsvScan(scan) => Ok(Arc::new(CsvExec::new(
- scan.base_conf.as_ref().unwrap().try_into()?,
- scan.has_header,
- str_to_byte(&scan.delimiter)?,
- ))),
- PhysicalPlanType::ParquetScan(scan) => {
- Ok(Arc::new(ParquetExec::new(
- scan.base_conf.as_ref().unwrap().try_into()?,
- // TODO predicate should be de-serialized
- None,
- )))
- }
- PhysicalPlanType::AvroScan(scan) => Ok(Arc::new(AvroExec::new(
- scan.base_conf.as_ref().unwrap().try_into()?,
- ))),
- PhysicalPlanType::CoalesceBatches(coalesce_batches) => {
- let input: Arc<dyn ExecutionPlan> =
- convert_box_required!(coalesce_batches.input)?;
- Ok(Arc::new(CoalesceBatchesExec::new(
- input,
- coalesce_batches.target_batch_size as usize,
- )))
- }
- PhysicalPlanType::Merge(merge) => {
- let input: Arc<dyn ExecutionPlan> = convert_box_required!(merge.input)?;
- Ok(Arc::new(CoalescePartitionsExec::new(input)))
- }
- PhysicalPlanType::Repartition(repart) => {
- let input: Arc<dyn ExecutionPlan> = convert_box_required!(repart.input)?;
- match repart.partition_method {
- Some(PartitionMethod::Hash(ref hash_part)) => {
- let expr = hash_part
- .hash_expr
- .iter()
- .map(|e| e.try_into())
- .collect::<Result<Vec<Arc<dyn PhysicalExpr>>, _>>()?;
-
- Ok(Arc::new(RepartitionExec::try_new(
- input,
- Partitioning::Hash(
- expr,
- hash_part.partition_count.try_into().unwrap(),
- ),
- )?))
- }
- Some(PartitionMethod::RoundRobin(partition_count)) => {
- Ok(Arc::new(RepartitionExec::try_new(
- input,
- Partitioning::RoundRobinBatch(
- partition_count.try_into().unwrap(),
- ),
- )?))
- }
- Some(PartitionMethod::Unknown(partition_count)) => {
- Ok(Arc::new(RepartitionExec::try_new(
- input,
- Partitioning::UnknownPartitioning(
- partition_count.try_into().unwrap(),
- ),
- )?))
- }
- _ => Err(BallistaError::General(
- "Invalid partitioning scheme".to_owned(),
- )),
- }
- }
- PhysicalPlanType::GlobalLimit(limit) => {
- let input: Arc<dyn ExecutionPlan> = convert_box_required!(limit.input)?;
- Ok(Arc::new(GlobalLimitExec::new(input, limit.limit as usize)))
- }
- PhysicalPlanType::LocalLimit(limit) => {
- let input: Arc<dyn ExecutionPlan> = convert_box_required!(limit.input)?;
- Ok(Arc::new(LocalLimitExec::new(input, limit.limit as usize)))
- }
- PhysicalPlanType::Window(window_agg) => {
- let input: Arc<dyn ExecutionPlan> =
- convert_box_required!(window_agg.input)?;
- let input_schema = window_agg
- .input_schema
- .as_ref()
- .ok_or_else(|| {
- BallistaError::General(
- "input_schema in WindowAggrNode is missing.".to_owned(),
- )
- })?
- .clone();
- let physical_schema: SchemaRef =
- SchemaRef::new((&input_schema).try_into()?);
-
- let physical_window_expr: Vec<Arc<dyn WindowExpr>> = window_agg
- .window_expr
- .iter()
- .zip(window_agg.window_expr_name.iter())
- .map(|(expr, name)| {
- let expr_type = expr.expr_type.as_ref().ok_or_else(|| {
- proto_error("Unexpected empty window physical expression")
- })?;
-
- match expr_type {
- ExprType::WindowExpr(window_node) => Ok(create_window_expr(
- &convert_required!(window_node.window_function)?,
- name.to_owned(),
- &[convert_box_required!(window_node.expr)?],
- &[],
- &[],
- Some(WindowFrame::default()),
- &physical_schema,
- )?),
- _ => Err(BallistaError::General(
- "Invalid expression for WindowAggrExec".to_string(),
- )),
- }
- })
- .collect::<Result<Vec<_>, _>>()?;
-
- Ok(Arc::new(WindowAggExec::try_new(
- physical_window_expr,
- input,
- Arc::new((&input_schema).try_into()?),
- )?))
- }
- PhysicalPlanType::HashAggregate(hash_agg) => {
- let input: Arc<dyn ExecutionPlan> =
- convert_box_required!(hash_agg.input)?;
- let mode = protobuf::AggregateMode::from_i32(hash_agg.mode).ok_or_else(|| {
- proto_error(format!(
- "Received a HashAggregateNode message with unknown AggregateMode {}",
- hash_agg.mode
- ))
- })?;
- let agg_mode: AggregateMode = match mode {
- protobuf::AggregateMode::Partial => AggregateMode::Partial,
- protobuf::AggregateMode::Final => AggregateMode::Final,
- protobuf::AggregateMode::FinalPartitioned => {
- AggregateMode::FinalPartitioned
- }
- };
- let group = hash_agg
- .group_expr
- .iter()
- .zip(hash_agg.group_expr_name.iter())
- .map(|(expr, name)| {
- expr.try_into().map(|expr| (expr, name.to_string()))
- })
- .collect::<Result<Vec<_>, _>>()?;
-
- let input_schema = hash_agg
- .input_schema
- .as_ref()
- .ok_or_else(|| {
- BallistaError::General(
- "input_schema in HashAggregateNode is missing.".to_owned(),
- )
- })?
- .clone();
- let physical_schema: SchemaRef =
- SchemaRef::new((&input_schema).try_into()?);
-
- let physical_aggr_expr: Vec<Arc<dyn AggregateExpr>> = hash_agg
- .aggr_expr
- .iter()
- .zip(hash_agg.aggr_expr_name.iter())
- .map(|(expr, name)| {
- let expr_type = expr.expr_type.as_ref().ok_or_else(|| {
- proto_error("Unexpected empty aggregate physical expression")
- })?;
-
- match expr_type {
- ExprType::AggregateExpr(agg_node) => {
- let aggr_function =
- protobuf::AggregateFunction::from_i32(
- agg_node.aggr_function,
- )
- .ok_or_else(
- || {
- proto_error(format!(
- "Received an unknown aggregate function: {}",
- agg_node.aggr_function
- ))
- },
- )?;
-
- Ok(create_aggregate_expr(
- &aggr_function.into(),
- false,
- &[convert_box_required!(agg_node.expr)?],
- &physical_schema,
- name.to_string(),
- )?)
- }
- _ => Err(BallistaError::General(
- "Invalid aggregate expression for HashAggregateExec"
- .to_string(),
- )),
- }
- })
- .collect::<Result<Vec<_>, _>>()?;
-
- Ok(Arc::new(HashAggregateExec::try_new(
- agg_mode,
- group,
- physical_aggr_expr,
- input,
- Arc::new((&input_schema).try_into()?),
- )?))
- }
- PhysicalPlanType::HashJoin(hashjoin) => {
- let left: Arc<dyn ExecutionPlan> = convert_box_required!(hashjoin.left)?;
- let right: Arc<dyn ExecutionPlan> =
- convert_box_required!(hashjoin.right)?;
- let on: Vec<(Column, Column)> = hashjoin
- .on
- .iter()
- .map(|col| {
- let left = into_required!(col.left)?;
- let right = into_required!(col.right)?;
- Ok((left, right))
- })
- .collect::<Result<_, Self::Error>>()?;
- let join_type = protobuf::JoinType::from_i32(hashjoin.join_type)
- .ok_or_else(|| {
- proto_error(format!(
- "Received a HashJoinNode message with unknown JoinType {}",
- hashjoin.join_type
- ))
- })?;
-
- let partition_mode =
- protobuf::PartitionMode::from_i32(hashjoin.partition_mode)
- .ok_or_else(|| {
- proto_error(format!(
- "Received a HashJoinNode message with unknown PartitionMode {}",
- hashjoin.partition_mode
- ))
- })?;
- let partition_mode = match partition_mode {
- protobuf::PartitionMode::CollectLeft => PartitionMode::CollectLeft,
- protobuf::PartitionMode::Partitioned => PartitionMode::Partitioned,
- };
- Ok(Arc::new(HashJoinExec::try_new(
- left,
- right,
- on,
- &join_type.into(),
- partition_mode,
- &hashjoin.null_equals_null,
- )?))
- }
- PhysicalPlanType::CrossJoin(crossjoin) => {
- let left: Arc<dyn ExecutionPlan> = convert_box_required!(crossjoin.left)?;
- let right: Arc<dyn ExecutionPlan> =
- convert_box_required!(crossjoin.right)?;
- Ok(Arc::new(CrossJoinExec::try_new(left, right)?))
- }
- PhysicalPlanType::ShuffleWriter(shuffle_writer) => {
- let input: Arc<dyn ExecutionPlan> =
- convert_box_required!(shuffle_writer.input)?;
-
- let output_partitioning = parse_protobuf_hash_partitioning(
- shuffle_writer.output_partitioning.as_ref(),
- )?;
-
- Ok(Arc::new(ShuffleWriterExec::try_new(
- shuffle_writer.job_id.clone(),
- shuffle_writer.stage_id as usize,
- input,
- "".to_string(), // this is intentional but hacky - the executor will fill this in
- output_partitioning,
- )?))
- }
- PhysicalPlanType::ShuffleReader(shuffle_reader) => {
- let schema = Arc::new(convert_required!(shuffle_reader.schema)?);
- let partition_location: Vec<Vec<PartitionLocation>> = shuffle_reader
- .partition
- .iter()
- .map(|p| {
- p.location
- .iter()
- .map(|l| l.clone().try_into())
- .collect::<Result<Vec<_>, _>>()
- })
- .collect::<Result<Vec<_>, BallistaError>>()?;
- let shuffle_reader =
- ShuffleReaderExec::try_new(partition_location, schema)?;
- Ok(Arc::new(shuffle_reader))
- }
- PhysicalPlanType::Empty(empty) => {
- let schema = Arc::new(convert_required!(empty.schema)?);
- Ok(Arc::new(EmptyExec::new(empty.produce_one_row, schema)))
- }
- PhysicalPlanType::Sort(sort) => {
- let input: Arc<dyn ExecutionPlan> = convert_box_required!(sort.input)?;
- let exprs = sort
- .expr
- .iter()
- .map(|expr| {
- let expr = expr.expr_type.as_ref().ok_or_else(|| {
- proto_error(format!(
- "physical_plan::from_proto() Unexpected expr {:?}",
- self
- ))
- })?;
- if let protobuf::physical_expr_node::ExprType::Sort(sort_expr) = expr {
- let expr = sort_expr
- .expr
- .as_ref()
- .ok_or_else(|| {
- proto_error(format!(
- "physical_plan::from_proto() Unexpected sort expr {:?}",
- self
- ))
- })?
- .as_ref();
- Ok(PhysicalSortExpr {
- expr: expr.try_into()?,
- options: SortOptions {
- descending: !sort_expr.asc,
- nulls_first: sort_expr.nulls_first,
- },
- })
- } else {
- Err(BallistaError::General(format!(
- "physical_plan::from_proto() {:?}",
- self
- )))
- }
- })
- .collect::<Result<Vec<_>, _>>()?;
- Ok(Arc::new(SortExec::try_new(exprs, input)?))
- }
- PhysicalPlanType::Unresolved(unresolved_shuffle) => {
- let schema = Arc::new(convert_required!(unresolved_shuffle.schema)?);
- Ok(Arc::new(UnresolvedShuffleExec {
- stage_id: unresolved_shuffle.stage_id as usize,
- schema,
- input_partition_count: unresolved_shuffle.input_partition_count
- as usize,
- output_partition_count: unresolved_shuffle.output_partition_count
- as usize,
- }))
- }
- }
- }
-}
-
impl From<&protobuf::PhysicalColumn> for Column {
fn from(c: &protobuf::PhysicalColumn) -> Column {
Column::new(&c.name, c.index as usize)
diff --git a/ballista/rust/core/src/serde/physical_plan/mod.rs b/ballista/rust/core/src/serde/physical_plan/mod.rs
index a4f2f2f..50aa317 100644
--- a/ballista/rust/core/src/serde/physical_plan/mod.rs
+++ b/ballista/rust/core/src/serde/physical_plan/mod.rs
@@ -15,14 +15,922 @@
// specific language governing permissions and limitations
// under the License.
+use crate::error::BallistaError;
+use crate::execution_plans::{
+ ShuffleReaderExec, ShuffleWriterExec, UnresolvedShuffleExec,
+};
+use crate::serde::physical_plan::from_proto::parse_protobuf_hash_partitioning;
+use crate::serde::protobuf::physical_expr_node::ExprType;
+use crate::serde::protobuf::physical_plan_node::PhysicalPlanType;
+use crate::serde::protobuf::repartition_exec_node::PartitionMethod;
+use crate::serde::protobuf::ShuffleReaderPartition;
+use crate::serde::protobuf::{PhysicalExtensionNode, PhysicalPlanNode};
+use crate::serde::scheduler::PartitionLocation;
+use crate::serde::{
+ byte_to_string, proto_error, protobuf, str_to_byte, AsExecutionPlan,
+ PhysicalExtensionCodec,
+};
+use crate::{convert_box_required, convert_required, into_physical_plan, into_required};
+use datafusion::arrow::compute::SortOptions;
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::datasource::object_store::local::LocalFileSystem;
+use datafusion::datasource::PartitionedFile;
+use datafusion::logical_plan::window_frames::WindowFrame;
+use datafusion::physical_plan::aggregates::create_aggregate_expr;
+use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
+use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion::physical_plan::cross_join::CrossJoinExec;
+use datafusion::physical_plan::empty::EmptyExec;
+use datafusion::physical_plan::expressions::{Column, PhysicalSortExpr};
+use datafusion::physical_plan::file_format::{
+ AvroExec, CsvExec, FileScanConfig, ParquetExec,
+};
+use datafusion::physical_plan::filter::FilterExec;
+use datafusion::physical_plan::hash_aggregate::{AggregateMode, HashAggregateExec};
+use datafusion::physical_plan::hash_join::{HashJoinExec, PartitionMode};
+use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
+use datafusion::physical_plan::projection::ProjectionExec;
+use datafusion::physical_plan::repartition::RepartitionExec;
+use datafusion::physical_plan::sorts::sort::SortExec;
+use datafusion::physical_plan::windows::{create_window_expr, WindowAggExec};
+use datafusion::physical_plan::{
+ AggregateExpr, ExecutionPlan, Partitioning, PhysicalExpr, WindowExpr,
+};
+use datafusion::prelude::ExecutionContext;
+use prost::bytes::BufMut;
+use prost::Message;
+use std::convert::TryInto;
+use std::sync::Arc;
+
pub mod from_proto;
pub mod to_proto;
+impl AsExecutionPlan for PhysicalPlanNode {
+ fn try_decode(buf: &[u8]) -> Result<Self, BallistaError>
+ where
+ Self: Sized,
+ {
+ PhysicalPlanNode::decode(buf).map_err(|e| {
+ BallistaError::Internal(format!("failed to decode physical plan: {:?}", e))
+ })
+ }
+
+ fn try_encode<B>(&self, buf: &mut B) -> Result<(), BallistaError>
+ where
+ B: BufMut,
+ Self: Sized,
+ {
+ self.encode(buf).map_err(|e| {
+ BallistaError::Internal(format!("failed to encode physical plan: {:?}", e))
+ })
+ }
+
+ fn try_into_physical_plan(
+ &self,
+ ctx: &ExecutionContext,
+ extension_codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<Arc<dyn ExecutionPlan>, BallistaError> {
+ let plan = self.physical_plan_type.as_ref().ok_or_else(|| {
+ proto_error(format!(
+ "physical_plan::from_proto() Unsupported physical plan '{:?}'",
+ self
+ ))
+ })?;
+ match plan {
+ PhysicalPlanType::Projection(projection) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(projection.input, ctx, extension_codec)?;
+ let exprs = projection
+ .expr
+ .iter()
+ .zip(projection.expr_name.iter())
+ .map(|(expr, name)| Ok((expr.try_into()?, name.to_string())))
+ .collect::<Result<Vec<(Arc<dyn PhysicalExpr>, String)>, BallistaError>>(
+ )?;
+ Ok(Arc::new(ProjectionExec::try_new(exprs, input)?))
+ }
+ PhysicalPlanType::Filter(filter) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(filter.input, ctx, extension_codec)?;
+ let predicate = filter
+ .expr
+ .as_ref()
+ .ok_or_else(|| {
+ BallistaError::General(
+ "filter (FilterExecNode) in PhysicalPlanNode is missing."
+ .to_owned(),
+ )
+ })?
+ .try_into()?;
+ Ok(Arc::new(FilterExec::try_new(predicate, input)?))
+ }
+ PhysicalPlanType::CsvScan(scan) => Ok(Arc::new(CsvExec::new(
+ decode_scan_config(scan.base_conf.as_ref().unwrap(), ctx)?,
+ scan.has_header,
+ str_to_byte(&scan.delimiter)?,
+ ))),
+ PhysicalPlanType::ParquetScan(scan) => {
+ Ok(Arc::new(ParquetExec::new(
+ decode_scan_config(scan.base_conf.as_ref().unwrap(), ctx)?,
+ // TODO predicate should be de-serialized
+ None,
+ )))
+ }
+ PhysicalPlanType::AvroScan(scan) => Ok(Arc::new(AvroExec::new(
+ decode_scan_config(scan.base_conf.as_ref().unwrap(), ctx)?,
+ ))),
+ PhysicalPlanType::CoalesceBatches(coalesce_batches) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(coalesce_batches.input, ctx, extension_codec)?;
+ Ok(Arc::new(CoalesceBatchesExec::new(
+ input,
+ coalesce_batches.target_batch_size as usize,
+ )))
+ }
+ PhysicalPlanType::Merge(merge) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(merge.input, ctx, extension_codec)?;
+ Ok(Arc::new(CoalescePartitionsExec::new(input)))
+ }
+ PhysicalPlanType::Repartition(repart) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(repart.input, ctx, extension_codec)?;
+ match repart.partition_method {
+ Some(PartitionMethod::Hash(ref hash_part)) => {
+ let expr = hash_part
+ .hash_expr
+ .iter()
+ .map(|e| e.try_into())
+ .collect::<Result<Vec<Arc<dyn PhysicalExpr>>, _>>()?;
+
+ Ok(Arc::new(RepartitionExec::try_new(
+ input,
+ Partitioning::Hash(
+ expr,
+ hash_part.partition_count.try_into().unwrap(),
+ ),
+ )?))
+ }
+ Some(PartitionMethod::RoundRobin(partition_count)) => {
+ Ok(Arc::new(RepartitionExec::try_new(
+ input,
+ Partitioning::RoundRobinBatch(
+ partition_count.try_into().unwrap(),
+ ),
+ )?))
+ }
+ Some(PartitionMethod::Unknown(partition_count)) => {
+ Ok(Arc::new(RepartitionExec::try_new(
+ input,
+ Partitioning::UnknownPartitioning(
+ partition_count.try_into().unwrap(),
+ ),
+ )?))
+ }
+ _ => Err(BallistaError::General(
+ "Invalid partitioning scheme".to_owned(),
+ )),
+ }
+ }
+ PhysicalPlanType::GlobalLimit(limit) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(limit.input, ctx, extension_codec)?;
+ Ok(Arc::new(GlobalLimitExec::new(input, limit.limit as usize)))
+ }
+ PhysicalPlanType::LocalLimit(limit) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(limit.input, ctx, extension_codec)?;
+ Ok(Arc::new(LocalLimitExec::new(input, limit.limit as usize)))
+ }
+ PhysicalPlanType::Window(window_agg) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(window_agg.input, ctx, extension_codec)?;
+ let input_schema = window_agg
+ .input_schema
+ .as_ref()
+ .ok_or_else(|| {
+ BallistaError::General(
+ "input_schema in WindowAggrNode is missing.".to_owned(),
+ )
+ })?
+ .clone();
+ let physical_schema: SchemaRef =
+ SchemaRef::new((&input_schema).try_into()?);
+
+ let physical_window_expr: Vec<Arc<dyn WindowExpr>> = window_agg
+ .window_expr
+ .iter()
+ .zip(window_agg.window_expr_name.iter())
+ .map(|(expr, name)| {
+ let expr_type = expr.expr_type.as_ref().ok_or_else(|| {
+ proto_error("Unexpected empty window physical expression")
+ })?;
+
+ match expr_type {
+ ExprType::WindowExpr(window_node) => Ok(create_window_expr(
+ &convert_required!(window_node.window_function)?,
+ name.to_owned(),
+ &[convert_box_required!(window_node.expr)?],
+ &[],
+ &[],
+ Some(WindowFrame::default()),
+ &physical_schema,
+ )?),
+ _ => Err(BallistaError::General(
+ "Invalid expression for WindowAggrExec".to_string(),
+ )),
+ }
+ })
+ .collect::<Result<Vec<_>, _>>()?;
+
+ Ok(Arc::new(WindowAggExec::try_new(
+ physical_window_expr,
+ input,
+ Arc::new((&input_schema).try_into()?),
+ )?))
+ }
+ PhysicalPlanType::HashAggregate(hash_agg) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(hash_agg.input, ctx, extension_codec)?;
+ let mode = protobuf::AggregateMode::from_i32(hash_agg.mode).ok_or_else(|| {
+ proto_error(format!(
+ "Received a HashAggregateNode message with unknown AggregateMode {}",
+ hash_agg.mode
+ ))
+ })?;
+ let agg_mode: AggregateMode = match mode {
+ protobuf::AggregateMode::Partial => AggregateMode::Partial,
+ protobuf::AggregateMode::Final => AggregateMode::Final,
+ protobuf::AggregateMode::FinalPartitioned => {
+ AggregateMode::FinalPartitioned
+ }
+ };
+ let group = hash_agg
+ .group_expr
+ .iter()
+ .zip(hash_agg.group_expr_name.iter())
+ .map(|(expr, name)| {
+ expr.try_into().map(|expr| (expr, name.to_string()))
+ })
+ .collect::<Result<Vec<_>, _>>()?;
+
+ let input_schema = hash_agg
+ .input_schema
+ .as_ref()
+ .ok_or_else(|| {
+ BallistaError::General(
+ "input_schema in HashAggregateNode is missing.".to_owned(),
+ )
+ })?
+ .clone();
+ let physical_schema: SchemaRef =
+ SchemaRef::new((&input_schema).try_into()?);
+
+ let physical_aggr_expr: Vec<Arc<dyn AggregateExpr>> = hash_agg
+ .aggr_expr
+ .iter()
+ .zip(hash_agg.aggr_expr_name.iter())
+ .map(|(expr, name)| {
+ let expr_type = expr.expr_type.as_ref().ok_or_else(|| {
+ proto_error("Unexpected empty aggregate physical expression")
+ })?;
+
+ match expr_type {
+ ExprType::AggregateExpr(agg_node) => {
+ let aggr_function =
+ protobuf::AggregateFunction::from_i32(
+ agg_node.aggr_function,
+ )
+ .ok_or_else(
+ || {
+ proto_error(format!(
+ "Received an unknown aggregate function: {}",
+ agg_node.aggr_function
+ ))
+ },
+ )?;
+
+ Ok(create_aggregate_expr(
+ &aggr_function.into(),
+ false,
+ &[convert_box_required!(agg_node.expr)?],
+ &physical_schema,
+ name.to_string(),
+ )?)
+ }
+ _ => Err(BallistaError::General(
+ "Invalid aggregate expression for HashAggregateExec"
+ .to_string(),
+ )),
+ }
+ })
+ .collect::<Result<Vec<_>, _>>()?;
+
+ Ok(Arc::new(HashAggregateExec::try_new(
+ agg_mode,
+ group,
+ physical_aggr_expr,
+ input,
+ Arc::new((&input_schema).try_into()?),
+ )?))
+ }
+ PhysicalPlanType::HashJoin(hashjoin) => {
+ let left: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(hashjoin.left, ctx, extension_codec)?;
+ let right: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(hashjoin.right, ctx, extension_codec)?;
+ let on: Vec<(Column, Column)> = hashjoin
+ .on
+ .iter()
+ .map(|col| {
+ let left = into_required!(col.left)?;
+ let right = into_required!(col.right)?;
+ Ok((left, right))
+ })
+ .collect::<Result<_, BallistaError>>()?;
+ let join_type = protobuf::JoinType::from_i32(hashjoin.join_type)
+ .ok_or_else(|| {
+ proto_error(format!(
+ "Received a HashJoinNode message with unknown JoinType {}",
+ hashjoin.join_type
+ ))
+ })?;
+
+ let partition_mode =
+ protobuf::PartitionMode::from_i32(hashjoin.partition_mode)
+ .ok_or_else(|| {
+ proto_error(format!(
+ "Received a HashJoinNode message with unknown PartitionMode {}",
+ hashjoin.partition_mode
+ ))
+ })?;
+ let partition_mode = match partition_mode {
+ protobuf::PartitionMode::CollectLeft => PartitionMode::CollectLeft,
+ protobuf::PartitionMode::Partitioned => PartitionMode::Partitioned,
+ };
+ Ok(Arc::new(HashJoinExec::try_new(
+ left,
+ right,
+ on,
+ &join_type.into(),
+ partition_mode,
+ &hashjoin.null_equals_null,
+ )?))
+ }
+ PhysicalPlanType::CrossJoin(crossjoin) => {
+ let left: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(crossjoin.left, ctx, extension_codec)?;
+ let right: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(crossjoin.right, ctx, extension_codec)?;
+ Ok(Arc::new(CrossJoinExec::try_new(left, right)?))
+ }
+ PhysicalPlanType::ShuffleWriter(shuffle_writer) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(shuffle_writer.input, ctx, extension_codec)?;
+
+ let output_partitioning = parse_protobuf_hash_partitioning(
+ shuffle_writer.output_partitioning.as_ref(),
+ )?;
+
+ Ok(Arc::new(ShuffleWriterExec::try_new(
+ shuffle_writer.job_id.clone(),
+ shuffle_writer.stage_id as usize,
+ input,
+ "".to_string(), // this is intentional but hacky - the executor will fill this in
+ output_partitioning,
+ )?))
+ }
+ PhysicalPlanType::ShuffleReader(shuffle_reader) => {
+ let schema = Arc::new(convert_required!(shuffle_reader.schema)?);
+ let partition_location: Vec<Vec<PartitionLocation>> = shuffle_reader
+ .partition
+ .iter()
+ .map(|p| {
+ p.location
+ .iter()
+ .map(|l| l.clone().try_into())
+ .collect::<Result<Vec<_>, _>>()
+ })
+ .collect::<Result<Vec<_>, BallistaError>>()?;
+ let shuffle_reader =
+ ShuffleReaderExec::try_new(partition_location, schema)?;
+ Ok(Arc::new(shuffle_reader))
+ }
+ PhysicalPlanType::Empty(empty) => {
+ let schema = Arc::new(convert_required!(empty.schema)?);
+ Ok(Arc::new(EmptyExec::new(empty.produce_one_row, schema)))
+ }
+ PhysicalPlanType::Sort(sort) => {
+ let input: Arc<dyn ExecutionPlan> =
+ into_physical_plan!(sort.input, ctx, extension_codec)?;
+ let exprs = sort
+ .expr
+ .iter()
+ .map(|expr| {
+ let expr = expr.expr_type.as_ref().ok_or_else(|| {
+ proto_error(format!(
+ "physical_plan::from_proto() Unexpected expr {:?}",
+ self
+ ))
+ })?;
+ if let protobuf::physical_expr_node::ExprType::Sort(sort_expr) = expr {
+ let expr = sort_expr
+ .expr
+ .as_ref()
+ .ok_or_else(|| {
+ proto_error(format!(
+ "physical_plan::from_proto() Unexpected sort expr {:?}",
+ self
+ ))
+ })?
+ .as_ref();
+ Ok(PhysicalSortExpr {
+ expr: expr.try_into()?,
+ options: SortOptions {
+ descending: !sort_expr.asc,
+ nulls_first: sort_expr.nulls_first,
+ },
+ })
+ } else {
+ Err(BallistaError::General(format!(
+ "physical_plan::from_proto() {:?}",
+ self
+ )))
+ }
+ })
+ .collect::<Result<Vec<_>, _>>()?;
+ Ok(Arc::new(SortExec::try_new(exprs, input)?))
+ }
+ PhysicalPlanType::Unresolved(unresolved_shuffle) => {
+ let schema = Arc::new(convert_required!(unresolved_shuffle.schema)?);
+ Ok(Arc::new(UnresolvedShuffleExec {
+ stage_id: unresolved_shuffle.stage_id as usize,
+ schema,
+ input_partition_count: unresolved_shuffle.input_partition_count
+ as usize,
+ output_partition_count: unresolved_shuffle.output_partition_count
+ as usize,
+ }))
+ }
+ PhysicalPlanType::Extension(extension) => {
+ let inputs: Vec<Arc<dyn ExecutionPlan>> = extension
+ .inputs
+ .iter()
+ .map(|i| i.try_into_physical_plan(ctx, extension_codec))
+ .collect::<Result<_, BallistaError>>()?;
+
+ let extension_node =
+ extension_codec.try_decode(extension.node.as_slice(), &inputs)?;
+
+ Ok(extension_node)
+ }
+ }
+ }
+
+ fn try_from_physical_plan(
+ plan: Arc<dyn ExecutionPlan>,
+ extension_codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<Self, BallistaError>
+ where
+ Self: Sized,
+ {
+ let plan_clone = plan.clone();
+ let plan = plan.as_any();
+
+ if let Some(exec) = plan.downcast_ref::<ProjectionExec>() {
+ let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.input().to_owned(),
+ extension_codec,
+ )?;
+ let expr = exec
+ .expr()
+ .iter()
+ .map(|expr| expr.0.clone().try_into())
+ .collect::<Result<Vec<_>, BallistaError>>()?;
+ let expr_name = exec.expr().iter().map(|expr| expr.1.clone()).collect();
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::Projection(Box::new(
+ protobuf::ProjectionExecNode {
+ input: Some(Box::new(input)),
+ expr,
+ expr_name,
+ },
+ ))),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<FilterExec>() {
+ let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.input().to_owned(),
+ extension_codec,
+ )?;
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::Filter(Box::new(
+ protobuf::FilterExecNode {
+ input: Some(Box::new(input)),
+ expr: Some(exec.predicate().clone().try_into()?),
+ },
+ ))),
+ })
+ } else if let Some(limit) = plan.downcast_ref::<GlobalLimitExec>() {
+ let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ limit.input().to_owned(),
+ extension_codec,
+ )?;
+
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::GlobalLimit(Box::new(
+ protobuf::GlobalLimitExecNode {
+ input: Some(Box::new(input)),
+ limit: limit.limit() as u32,
+ },
+ ))),
+ })
+ } else if let Some(limit) = plan.downcast_ref::<LocalLimitExec>() {
+ let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ limit.input().to_owned(),
+ extension_codec,
+ )?;
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::LocalLimit(Box::new(
+ protobuf::LocalLimitExecNode {
+ input: Some(Box::new(input)),
+ limit: limit.limit() as u32,
+ },
+ ))),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<HashJoinExec>() {
+ let left = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.left().to_owned(),
+ extension_codec,
+ )?;
+ let right = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.right().to_owned(),
+ extension_codec,
+ )?;
+ let on: Vec<protobuf::JoinOn> = exec
+ .on()
+ .iter()
+ .map(|tuple| protobuf::JoinOn {
+ left: Some(protobuf::PhysicalColumn {
+ name: tuple.0.name().to_string(),
+ index: tuple.0.index() as u32,
+ }),
+ right: Some(protobuf::PhysicalColumn {
+ name: tuple.1.name().to_string(),
+ index: tuple.1.index() as u32,
+ }),
+ })
+ .collect();
+ let join_type: protobuf::JoinType = exec.join_type().to_owned().into();
+
+ let partition_mode = match exec.partition_mode() {
+ PartitionMode::CollectLeft => protobuf::PartitionMode::CollectLeft,
+ PartitionMode::Partitioned => protobuf::PartitionMode::Partitioned,
+ };
+
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::HashJoin(Box::new(
+ protobuf::HashJoinExecNode {
+ left: Some(Box::new(left)),
+ right: Some(Box::new(right)),
+ on,
+ join_type: join_type.into(),
+ partition_mode: partition_mode.into(),
+ null_equals_null: *exec.null_equals_null(),
+ },
+ ))),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<CrossJoinExec>() {
+ let left = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.left().to_owned(),
+ extension_codec,
+ )?;
+ let right = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.right().to_owned(),
+ extension_codec,
+ )?;
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::CrossJoin(Box::new(
+ protobuf::CrossJoinExecNode {
+ left: Some(Box::new(left)),
+ right: Some(Box::new(right)),
+ },
+ ))),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<HashAggregateExec>() {
+ let groups = exec
+ .group_expr()
+ .iter()
+ .map(|expr| expr.0.to_owned().try_into())
+ .collect::<Result<Vec<_>, BallistaError>>()?;
+ let group_names = exec
+ .group_expr()
+ .iter()
+ .map(|expr| expr.1.to_owned())
+ .collect();
+ let agg = exec
+ .aggr_expr()
+ .iter()
+ .map(|expr| expr.to_owned().try_into())
+ .collect::<Result<Vec<_>, BallistaError>>()?;
+ let agg_names = exec
+ .aggr_expr()
+ .iter()
+ .map(|expr| match expr.field() {
+ Ok(field) => Ok(field.name().clone()),
+ Err(e) => Err(BallistaError::DataFusionError(e)),
+ })
+ .collect::<Result<_, BallistaError>>()?;
+
+ let agg_mode = match exec.mode() {
+ AggregateMode::Partial => protobuf::AggregateMode::Partial,
+ AggregateMode::Final => protobuf::AggregateMode::Final,
+ AggregateMode::FinalPartitioned => {
+ protobuf::AggregateMode::FinalPartitioned
+ }
+ };
+ let input_schema = exec.input_schema();
+ let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.input().to_owned(),
+ extension_codec,
+ )?;
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::HashAggregate(Box::new(
+ protobuf::HashAggregateExecNode {
+ group_expr: groups,
+ group_expr_name: group_names,
+ aggr_expr: agg,
+ aggr_expr_name: agg_names,
+ mode: agg_mode as i32,
+ input: Some(Box::new(input)),
+ input_schema: Some(input_schema.as_ref().into()),
+ },
+ ))),
+ })
+ } else if let Some(empty) = plan.downcast_ref::<EmptyExec>() {
+ let schema = empty.schema().as_ref().into();
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::Empty(
+ protobuf::EmptyExecNode {
+ produce_one_row: empty.produce_one_row(),
+ schema: Some(schema),
+ },
+ )),
+ })
+ } else if let Some(coalesce_batches) = plan.downcast_ref::<CoalesceBatchesExec>()
+ {
+ let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ coalesce_batches.input().to_owned(),
+ extension_codec,
+ )?;
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::CoalesceBatches(Box::new(
+ protobuf::CoalesceBatchesExecNode {
+ input: Some(Box::new(input)),
+ target_batch_size: coalesce_batches.target_batch_size() as u32,
+ },
+ ))),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<CsvExec>() {
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::CsvScan(
+ protobuf::CsvScanExecNode {
+ base_conf: Some(exec.base_config().try_into()?),
+ has_header: exec.has_header(),
+ delimiter: byte_to_string(exec.delimiter())?,
+ },
+ )),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<ParquetExec>() {
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::ParquetScan(
+ protobuf::ParquetScanExecNode {
+ base_conf: Some(exec.base_config().try_into()?),
+ // TODO serialize predicates
+ },
+ )),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<AvroExec>() {
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::AvroScan(
+ protobuf::AvroScanExecNode {
+ base_conf: Some(exec.base_config().try_into()?),
+ },
+ )),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<ShuffleReaderExec>() {
+ let mut partition = vec![];
+ for location in &exec.partition {
+ partition.push(protobuf::ShuffleReaderPartition {
+ location: location
+ .iter()
+ .map(|l| l.clone().try_into())
+ .collect::<Result<Vec<_>, _>>()?,
+ });
+ }
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::ShuffleReader(
+ protobuf::ShuffleReaderExecNode {
+ partition,
+ schema: Some(exec.schema().as_ref().into()),
+ },
+ )),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<CoalescePartitionsExec>() {
+ let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.input().to_owned(),
+ extension_codec,
+ )?;
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::Merge(Box::new(
+ protobuf::CoalescePartitionsExecNode {
+ input: Some(Box::new(input)),
+ },
+ ))),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<RepartitionExec>() {
+ let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.input().to_owned(),
+ extension_codec,
+ )?;
+
+ let pb_partition_method = match exec.partitioning() {
+ Partitioning::Hash(exprs, partition_count) => {
+ PartitionMethod::Hash(protobuf::PhysicalHashRepartition {
+ hash_expr: exprs
+ .iter()
+ .map(|expr| expr.clone().try_into())
+ .collect::<Result<Vec<_>, BallistaError>>()?,
+ partition_count: *partition_count as u64,
+ })
+ }
+ Partitioning::RoundRobinBatch(partition_count) => {
+ PartitionMethod::RoundRobin(*partition_count as u64)
+ }
+ Partitioning::UnknownPartitioning(partition_count) => {
+ PartitionMethod::Unknown(*partition_count as u64)
+ }
+ };
+
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::Repartition(Box::new(
+ protobuf::RepartitionExecNode {
+ input: Some(Box::new(input)),
+ partition_method: Some(pb_partition_method),
+ },
+ ))),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<SortExec>() {
+ let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.input().to_owned(),
+ extension_codec,
+ )?;
+ let expr = exec
+ .expr()
+ .iter()
+ .map(|expr| {
+ let sort_expr = Box::new(protobuf::PhysicalSortExprNode {
+ expr: Some(Box::new(expr.expr.to_owned().try_into()?)),
+ asc: !expr.options.descending,
+ nulls_first: expr.options.nulls_first,
+ });
+ Ok(protobuf::PhysicalExprNode {
+ expr_type: Some(protobuf::physical_expr_node::ExprType::Sort(
+ sort_expr,
+ )),
+ })
+ })
+ .collect::<Result<Vec<_>, BallistaError>>()?;
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::Sort(Box::new(
+ protobuf::SortExecNode {
+ input: Some(Box::new(input)),
+ expr,
+ },
+ ))),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<ShuffleWriterExec>() {
+ let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec.children()[0].to_owned(),
+ extension_codec,
+ )?;
+ // note that we use shuffle_output_partitioning() rather than output_partitioning()
+ // to get the true output partitioning
+ let output_partitioning = match exec.shuffle_output_partitioning() {
+ Some(Partitioning::Hash(exprs, partition_count)) => {
+ Some(protobuf::PhysicalHashRepartition {
+ hash_expr: exprs
+ .iter()
+ .map(|expr| expr.clone().try_into())
+ .collect::<Result<Vec<_>, BallistaError>>()?,
+ partition_count: *partition_count as u64,
+ })
+ }
+ None => None,
+ other => {
+ return Err(BallistaError::General(format!(
+ "physical_plan::to_proto() invalid partitioning for ShuffleWriterExec: {:?}",
+ other
+ )))
+ }
+ };
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::ShuffleWriter(Box::new(
+ protobuf::ShuffleWriterExecNode {
+ job_id: exec.job_id().to_string(),
+ stage_id: exec.stage_id() as u32,
+ input: Some(Box::new(input)),
+ output_partitioning,
+ },
+ ))),
+ })
+ } else if let Some(exec) = plan.downcast_ref::<UnresolvedShuffleExec>() {
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::Unresolved(
+ protobuf::UnresolvedShuffleExecNode {
+ stage_id: exec.stage_id as u32,
+ schema: Some(exec.schema().as_ref().into()),
+ input_partition_count: exec.input_partition_count as u32,
+ output_partition_count: exec.output_partition_count as u32,
+ },
+ )),
+ })
+ } else {
+ let mut buf: Vec<u8> = vec![];
+ extension_codec.try_encode(plan_clone.clone(), &mut buf)?;
+
+ let inputs: Vec<PhysicalPlanNode> = plan_clone
+ .children()
+ .into_iter()
+ .map(|i| PhysicalPlanNode::try_from_physical_plan(i, extension_codec))
+ .collect::<Result<_, BallistaError>>()?;
+
+ Ok(protobuf::PhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::Extension(
+ PhysicalExtensionNode { node: buf, inputs },
+ )),
+ })
+ }
+ }
+}
+
+fn decode_scan_config(
+ proto: &protobuf::FileScanExecConf,
+ ctx: &ExecutionContext,
+) -> Result<FileScanConfig, BallistaError> {
+ let schema = Arc::new(convert_required!(proto.schema)?);
+ let projection = proto
+ .projection
+ .iter()
+ .map(|i| *i as usize)
+ .collect::<Vec<_>>();
+ let projection = if projection.is_empty() {
+ None
+ } else {
+ Some(projection)
+ };
+ let statistics = convert_required!(proto.statistics)?;
+
+ let file_groups: Vec<Vec<PartitionedFile>> = proto
+ .file_groups
+ .iter()
+ .map(|f| f.try_into())
+ .collect::<Result<Vec<_>, _>>()?;
+
+ let object_store = if let Some(file) = file_groups.get(0).and_then(|h| h.get(0)) {
+ ctx.object_store(file.file_meta.path())?.0
+ } else {
+ Arc::new(LocalFileSystem {})
+ };
+
+ Ok(FileScanConfig {
+ object_store,
+ file_schema: schema,
+ file_groups,
+ statistics,
+ projection,
+ limit: proto.limit.as_ref().map(|sl| sl.limit as usize),
+ table_partition_cols: vec![],
+ })
+}
+
+#[macro_export]
+macro_rules! into_physical_plan {
+ ($PB:expr, $CTX:expr, $CODEC:expr) => {{
+ if let Some(field) = $PB.as_ref() {
+ field.as_ref().try_into_physical_plan(&$CTX, $CODEC)
+ } else {
+ Err(proto_error("Missing required field in protobuf"))
+ }
+ }};
+}
+
#[cfg(test)]
mod roundtrip_tests {
use std::{convert::TryInto, sync::Arc};
+ use crate::serde::{AsExecutionPlan, BallistaCodec};
use datafusion::physical_plan::sorts::sort::SortExec;
+ use datafusion::prelude::ExecutionContext;
use datafusion::{
arrow::{
compute::kernels::sort::SortOptions,
@@ -46,10 +954,21 @@ mod roundtrip_tests {
use super::super::super::error::Result;
use super::super::protobuf;
use crate::execution_plans::ShuffleWriterExec;
+ use crate::serde::protobuf::{LogicalPlanNode, PhysicalPlanNode};
fn roundtrip_test(exec_plan: Arc<dyn ExecutionPlan>) -> Result<()> {
- let proto: protobuf::PhysicalPlanNode = exec_plan.clone().try_into()?;
- let result_exec_plan: Arc<dyn ExecutionPlan> = (&proto).try_into()?;
+ let ctx = ExecutionContext::new();
+ let codec: BallistaCodec<LogicalPlanNode, PhysicalPlanNode> =
+ BallistaCodec::default();
+ let proto: protobuf::PhysicalPlanNode =
+ protobuf::PhysicalPlanNode::try_from_physical_plan(
+ exec_plan.clone(),
+ codec.physical_extension_codec(),
+ )
+ .expect("to proto");
+ let result_exec_plan: Arc<dyn ExecutionPlan> = proto
+ .try_into_physical_plan(&ctx, codec.physical_extension_codec())
+ .expect("from proto");
assert_eq!(
format!("{:?}", exec_plan),
format!("{:?}", result_exec_plan)
diff --git a/ballista/rust/core/src/serde/physical_plan/to_proto.rs b/ballista/rust/core/src/serde/physical_plan/to_proto.rs
index 4c6db9a..e6f4b59 100644
--- a/ballista/rust/core/src/serde/physical_plan/to_proto.rs
+++ b/ballista/rust/core/src/serde/physical_plan/to_proto.rs
@@ -75,339 +75,6 @@ use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion::physical_plan::functions::{BuiltinScalarFunction, ScalarFunctionExpr};
use datafusion::physical_plan::repartition::RepartitionExec;
-impl TryInto<protobuf::PhysicalPlanNode> for Arc<dyn ExecutionPlan> {
- type Error = BallistaError;
-
- fn try_into(self) -> Result<protobuf::PhysicalPlanNode, Self::Error> {
- let plan = self.as_any();
-
- if let Some(exec) = plan.downcast_ref::<ProjectionExec>() {
- let input: protobuf::PhysicalPlanNode = exec.input().to_owned().try_into()?;
- let expr = exec
- .expr()
- .iter()
- .map(|expr| expr.0.clone().try_into())
- .collect::<Result<Vec<_>, Self::Error>>()?;
- let expr_name = exec.expr().iter().map(|expr| expr.1.clone()).collect();
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::Projection(Box::new(
- protobuf::ProjectionExecNode {
- input: Some(Box::new(input)),
- expr,
- expr_name,
- },
- ))),
- })
- } else if let Some(exec) = plan.downcast_ref::<FilterExec>() {
- let input: protobuf::PhysicalPlanNode = exec.input().to_owned().try_into()?;
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::Filter(Box::new(
- protobuf::FilterExecNode {
- input: Some(Box::new(input)),
- expr: Some(exec.predicate().clone().try_into()?),
- },
- ))),
- })
- } else if let Some(limit) = plan.downcast_ref::<GlobalLimitExec>() {
- let input: protobuf::PhysicalPlanNode =
- limit.input().to_owned().try_into()?;
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::GlobalLimit(Box::new(
- protobuf::GlobalLimitExecNode {
- input: Some(Box::new(input)),
- limit: limit.limit() as u32,
- },
- ))),
- })
- } else if let Some(limit) = plan.downcast_ref::<LocalLimitExec>() {
- let input: protobuf::PhysicalPlanNode =
- limit.input().to_owned().try_into()?;
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::LocalLimit(Box::new(
- protobuf::LocalLimitExecNode {
- input: Some(Box::new(input)),
- limit: limit.limit() as u32,
- },
- ))),
- })
- } else if let Some(exec) = plan.downcast_ref::<HashJoinExec>() {
- let left: protobuf::PhysicalPlanNode = exec.left().to_owned().try_into()?;
- let right: protobuf::PhysicalPlanNode = exec.right().to_owned().try_into()?;
- let on: Vec<protobuf::JoinOn> = exec
- .on()
- .iter()
- .map(|tuple| protobuf::JoinOn {
- left: Some(protobuf::PhysicalColumn {
- name: tuple.0.name().to_string(),
- index: tuple.0.index() as u32,
- }),
- right: Some(protobuf::PhysicalColumn {
- name: tuple.1.name().to_string(),
- index: tuple.1.index() as u32,
- }),
- })
- .collect();
- let join_type: protobuf::JoinType = exec.join_type().to_owned().into();
-
- let partition_mode = match exec.partition_mode() {
- PartitionMode::CollectLeft => protobuf::PartitionMode::CollectLeft,
- PartitionMode::Partitioned => protobuf::PartitionMode::Partitioned,
- };
-
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::HashJoin(Box::new(
- protobuf::HashJoinExecNode {
- left: Some(Box::new(left)),
- right: Some(Box::new(right)),
- on,
- join_type: join_type.into(),
- partition_mode: partition_mode.into(),
- null_equals_null: *exec.null_equals_null(),
- },
- ))),
- })
- } else if let Some(exec) = plan.downcast_ref::<CrossJoinExec>() {
- let left: protobuf::PhysicalPlanNode = exec.left().to_owned().try_into()?;
- let right: protobuf::PhysicalPlanNode = exec.right().to_owned().try_into()?;
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::CrossJoin(Box::new(
- protobuf::CrossJoinExecNode {
- left: Some(Box::new(left)),
- right: Some(Box::new(right)),
- },
- ))),
- })
- } else if let Some(exec) = plan.downcast_ref::<HashAggregateExec>() {
- let groups = exec
- .group_expr()
- .iter()
- .map(|expr| expr.0.to_owned().try_into())
- .collect::<Result<Vec<_>, BallistaError>>()?;
- let group_names = exec
- .group_expr()
- .iter()
- .map(|expr| expr.1.to_owned())
- .collect();
- let agg = exec
- .aggr_expr()
- .iter()
- .map(|expr| expr.to_owned().try_into())
- .collect::<Result<Vec<_>, BallistaError>>()?;
- let agg_names = exec
- .aggr_expr()
- .iter()
- .map(|expr| match expr.field() {
- Ok(field) => Ok(field.name().clone()),
- Err(e) => Err(BallistaError::DataFusionError(e)),
- })
- .collect::<Result<_, Self::Error>>()?;
-
- let agg_mode = match exec.mode() {
- AggregateMode::Partial => protobuf::AggregateMode::Partial,
- AggregateMode::Final => protobuf::AggregateMode::Final,
- AggregateMode::FinalPartitioned => {
- protobuf::AggregateMode::FinalPartitioned
- }
- };
- let input_schema = exec.input_schema();
- let input: protobuf::PhysicalPlanNode = exec.input().to_owned().try_into()?;
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::HashAggregate(Box::new(
- protobuf::HashAggregateExecNode {
- group_expr: groups,
- group_expr_name: group_names,
- aggr_expr: agg,
- aggr_expr_name: agg_names,
- mode: agg_mode as i32,
- input: Some(Box::new(input)),
- input_schema: Some(input_schema.as_ref().into()),
- },
- ))),
- })
- } else if let Some(empty) = plan.downcast_ref::<EmptyExec>() {
- let schema = empty.schema().as_ref().into();
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::Empty(
- protobuf::EmptyExecNode {
- produce_one_row: empty.produce_one_row(),
- schema: Some(schema),
- },
- )),
- })
- } else if let Some(coalesce_batches) = plan.downcast_ref::<CoalesceBatchesExec>()
- {
- let input: protobuf::PhysicalPlanNode =
- coalesce_batches.input().to_owned().try_into()?;
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::CoalesceBatches(Box::new(
- protobuf::CoalesceBatchesExecNode {
- input: Some(Box::new(input)),
- target_batch_size: coalesce_batches.target_batch_size() as u32,
- },
- ))),
- })
- } else if let Some(exec) = plan.downcast_ref::<CsvExec>() {
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::CsvScan(
- protobuf::CsvScanExecNode {
- base_conf: Some(exec.base_config().try_into()?),
- has_header: exec.has_header(),
- delimiter: byte_to_string(exec.delimiter())?,
- },
- )),
- })
- } else if let Some(exec) = plan.downcast_ref::<ParquetExec>() {
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::ParquetScan(
- protobuf::ParquetScanExecNode {
- base_conf: Some(exec.base_config().try_into()?),
- // TODO serialize predicates
- },
- )),
- })
- } else if let Some(exec) = plan.downcast_ref::<AvroExec>() {
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::AvroScan(
- protobuf::AvroScanExecNode {
- base_conf: Some(exec.base_config().try_into()?),
- },
- )),
- })
- } else if let Some(exec) = plan.downcast_ref::<ShuffleReaderExec>() {
- let mut partition = vec![];
- for location in &exec.partition {
- partition.push(protobuf::ShuffleReaderPartition {
- location: location
- .iter()
- .map(|l| l.clone().try_into())
- .collect::<Result<Vec<_>, _>>()?,
- });
- }
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::ShuffleReader(
- protobuf::ShuffleReaderExecNode {
- partition,
- schema: Some(exec.schema().as_ref().into()),
- },
- )),
- })
- } else if let Some(exec) = plan.downcast_ref::<CoalescePartitionsExec>() {
- let input: protobuf::PhysicalPlanNode = exec.input().to_owned().try_into()?;
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::Merge(Box::new(
- protobuf::CoalescePartitionsExecNode {
- input: Some(Box::new(input)),
- },
- ))),
- })
- } else if let Some(exec) = plan.downcast_ref::<RepartitionExec>() {
- let input: protobuf::PhysicalPlanNode = exec.input().to_owned().try_into()?;
-
- let pb_partition_method = match exec.partitioning() {
- Partitioning::Hash(exprs, partition_count) => {
- PartitionMethod::Hash(protobuf::PhysicalHashRepartition {
- hash_expr: exprs
- .iter()
- .map(|expr| expr.clone().try_into())
- .collect::<Result<Vec<_>, BallistaError>>()?,
- partition_count: *partition_count as u64,
- })
- }
- Partitioning::RoundRobinBatch(partition_count) => {
- PartitionMethod::RoundRobin(*partition_count as u64)
- }
- Partitioning::UnknownPartitioning(partition_count) => {
- PartitionMethod::Unknown(*partition_count as u64)
- }
- };
-
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::Repartition(Box::new(
- protobuf::RepartitionExecNode {
- input: Some(Box::new(input)),
- partition_method: Some(pb_partition_method),
- },
- ))),
- })
- } else if let Some(exec) = plan.downcast_ref::<SortExec>() {
- let input: protobuf::PhysicalPlanNode = exec.input().to_owned().try_into()?;
- let expr = exec
- .expr()
- .iter()
- .map(|expr| {
- let sort_expr = Box::new(protobuf::PhysicalSortExprNode {
- expr: Some(Box::new(expr.expr.to_owned().try_into()?)),
- asc: !expr.options.descending,
- nulls_first: expr.options.nulls_first,
- });
- Ok(protobuf::PhysicalExprNode {
- expr_type: Some(protobuf::physical_expr_node::ExprType::Sort(
- sort_expr,
- )),
- })
- })
- .collect::<Result<Vec<_>, Self::Error>>()?;
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::Sort(Box::new(
- protobuf::SortExecNode {
- input: Some(Box::new(input)),
- expr,
- },
- ))),
- })
- } else if let Some(exec) = plan.downcast_ref::<ShuffleWriterExec>() {
- let input: protobuf::PhysicalPlanNode =
- exec.children()[0].to_owned().try_into()?;
- // note that we use shuffle_output_partitioning() rather than output_partitioning()
- // to get the true output partitioning
- let output_partitioning = match exec.shuffle_output_partitioning() {
- Some(Partitioning::Hash(exprs, partition_count)) => {
- Some(protobuf::PhysicalHashRepartition {
- hash_expr: exprs
- .iter()
- .map(|expr| expr.clone().try_into())
- .collect::<Result<Vec<_>, BallistaError>>()?,
- partition_count: *partition_count as u64,
- })
- }
- None => None,
- other => {
- return Err(BallistaError::General(format!(
- "physical_plan::to_proto() invalid partitioning for ShuffleWriterExec: {:?}",
- other
- )))
- }
- };
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::ShuffleWriter(Box::new(
- protobuf::ShuffleWriterExecNode {
- job_id: exec.job_id().to_string(),
- stage_id: exec.stage_id() as u32,
- input: Some(Box::new(input)),
- output_partitioning,
- },
- ))),
- })
- } else if let Some(exec) = plan.downcast_ref::<UnresolvedShuffleExec>() {
- Ok(protobuf::PhysicalPlanNode {
- physical_plan_type: Some(PhysicalPlanType::Unresolved(
- protobuf::UnresolvedShuffleExecNode {
- stage_id: exec.stage_id as u32,
- schema: Some(exec.schema().as_ref().into()),
- input_partition_count: exec.input_partition_count as u32,
- output_partition_count: exec.output_partition_count as u32,
- },
- )),
- })
- } else {
- Err(BallistaError::General(format!(
- "physical plan to_proto unsupported plan {:?}",
- self
- )))
- }
- }
-}
-
impl TryInto<protobuf::PhysicalExprNode> for Arc<dyn AggregateExpr> {
type Error = BallistaError;
diff --git a/ballista/rust/core/src/serde/scheduler/to_proto.rs b/ballista/rust/core/src/serde/scheduler/to_proto.rs
index bdc88d0..71a02d3 100644
--- a/ballista/rust/core/src/serde/scheduler/to_proto.rs
+++ b/ballista/rust/core/src/serde/scheduler/to_proto.rs
@@ -48,23 +48,6 @@ impl TryInto<protobuf::Action> for Action {
}
}
-impl TryInto<protobuf::ExecutePartition> for ExecutePartition {
- type Error = BallistaError;
-
- fn try_into(self) -> Result<protobuf::ExecutePartition, Self::Error> {
- Ok(protobuf::ExecutePartition {
- job_id: self.job_id,
- stage_id: self.stage_id as u32,
- partition_id: self.partition_id.iter().map(|n| *n as u32).collect(),
- plan: Some(self.plan.try_into()?),
- partition_location: vec![],
- output_partitioning: hash_partitioning_to_proto(
- self.output_partitioning.as_ref(),
- )?,
- })
- }
-}
-
#[allow(clippy::from_over_into)]
impl Into<protobuf::PartitionId> for PartitionId {
fn into(self) -> protobuf::PartitionId {
diff --git a/ballista/rust/core/src/utils.rs b/ballista/rust/core/src/utils.rs
index efe175f..fae3826 100644
--- a/ballista/rust/core/src/utils.rs
+++ b/ballista/rust/core/src/utils.rs
@@ -17,6 +17,7 @@
use std::collections::HashMap;
use std::io::{BufWriter, Write};
+use std::marker::PhantomData;
use std::ops::Deref;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
@@ -29,6 +30,7 @@ use crate::execution_plans::{
use crate::serde::scheduler::PartitionStats;
use crate::config::BallistaConfig;
+use crate::serde::{AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec};
use async_trait::async_trait;
use datafusion::arrow::datatypes::Schema;
use datafusion::arrow::error::Result as ArrowResult;
@@ -236,38 +238,68 @@ fn build_exec_plan_diagram(
/// Create a DataFusion context that uses the BallistaQueryPlanner to send logical plans
/// to a Ballista scheduler
-pub fn create_df_ctx_with_ballista_query_planner(
+pub fn create_df_ctx_with_ballista_query_planner<T: 'static + AsLogicalPlan>(
scheduler_host: &str,
scheduler_port: u16,
config: &BallistaConfig,
) -> ExecutionContext {
let scheduler_url = format!("http://{}:{}", scheduler_host, scheduler_port);
+ let planner: Arc<BallistaQueryPlanner<T>> =
+ Arc::new(BallistaQueryPlanner::new(scheduler_url, config.clone()));
let config = ExecutionConfig::new()
- .with_query_planner(Arc::new(BallistaQueryPlanner::new(
- scheduler_url,
- config.clone(),
- )))
+ .with_query_planner(planner)
.with_target_partitions(config.default_shuffle_partitions())
.with_information_schema(true);
ExecutionContext::with_config(config)
}
-pub struct BallistaQueryPlanner {
+pub struct BallistaQueryPlanner<T: AsLogicalPlan> {
scheduler_url: String,
config: BallistaConfig,
+ extension_codec: Arc<dyn LogicalExtensionCodec>,
+ plan_repr: PhantomData<T>,
}
-impl BallistaQueryPlanner {
+impl<T: 'static + AsLogicalPlan> BallistaQueryPlanner<T> {
pub fn new(scheduler_url: String, config: BallistaConfig) -> Self {
Self {
scheduler_url,
config,
+ extension_codec: Arc::new(DefaultLogicalExtensionCodec {}),
+ plan_repr: PhantomData,
+ }
+ }
+
+ pub fn with_extension(
+ scheduler_url: String,
+ config: BallistaConfig,
+ extension_codec: Arc<dyn LogicalExtensionCodec>,
+ ) -> Self {
+ Self {
+ scheduler_url,
+ config,
+ extension_codec,
+ plan_repr: PhantomData,
+ }
+ }
+
+ pub fn with_repr(
+ scheduler_url: String,
+ config: BallistaConfig,
+ extension_codec: Arc<dyn LogicalExtensionCodec>,
+ plan_repr: PhantomData<T>,
+ ) -> Self {
+ Self {
+ scheduler_url,
+ config,
+ extension_codec,
+ plan_repr,
}
}
}
#[async_trait]
-impl QueryPlanner for BallistaQueryPlanner {
+impl<T: 'static + AsLogicalPlan> QueryPlanner for BallistaQueryPlanner<T> {
async fn create_physical_plan(
&self,
logical_plan: &LogicalPlan,
@@ -278,10 +310,12 @@ impl QueryPlanner for BallistaQueryPlanner {
// table state is managed locally in the BallistaContext, not in the scheduler
Ok(Arc::new(EmptyExec::new(false, Arc::new(Schema::empty()))))
}
- _ => Ok(Arc::new(DistributedQueryExec::new(
+ _ => Ok(Arc::new(DistributedQueryExec::with_repr(
self.scheduler_url.clone(),
self.config.clone(),
logical_plan.clone(),
+ self.extension_codec.clone(),
+ self.plan_repr,
))),
}
}
diff --git a/ballista/rust/executor/src/execution_loop.rs b/ballista/rust/executor/src/execution_loop.rs
index 69bc838..a324ba4 100644
--- a/ballista/rust/executor/src/execution_loop.rs
+++ b/ballista/rust/executor/src/execution_loop.rs
@@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.
-use std::convert::TryInto;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{Receiver, Sender, TryRecvError};
use std::{sync::Arc, time::Duration};
@@ -34,12 +33,14 @@ use crate::as_task_status;
use crate::executor::Executor;
use ballista_core::error::BallistaError;
use ballista_core::serde::physical_plan::from_proto::parse_protobuf_hash_partitioning;
+use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec};
-pub async fn poll_loop(
+pub async fn poll_loop<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
mut scheduler: SchedulerGrpcClient<Channel>,
executor: Arc<Executor>,
executor_meta: ExecutorRegistration,
concurrent_tasks: usize,
+ codec: BallistaCodec<T, U>,
) {
let available_tasks_slots = Arc::new(AtomicUsize::new(concurrent_tasks));
let (task_status_sender, mut task_status_receiver) =
@@ -77,6 +78,7 @@ pub async fn poll_loop(
available_tasks_slots.clone(),
task_status_sender,
task,
+ &codec,
)
.await
{
@@ -102,12 +104,13 @@ pub async fn poll_loop(
}
}
-async fn run_received_tasks(
+async fn run_received_tasks<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
executor: Arc<Executor>,
executor_id: String,
available_tasks_slots: Arc<AtomicUsize>,
task_status_sender: Sender<TaskStatus>,
task: TaskDefinition,
+ codec: &BallistaCodec<T, U>,
) -> Result<(), BallistaError> {
let task_id = task.task_id.unwrap();
let task_id_log = format!(
@@ -116,7 +119,15 @@ async fn run_received_tasks(
);
info!("Received task {}", task_id_log);
available_tasks_slots.fetch_sub(1, Ordering::SeqCst);
- let plan: Arc<dyn ExecutionPlan> = (&task.plan.unwrap()).try_into().unwrap();
+
+ let plan: Arc<dyn ExecutionPlan> =
+ U::try_decode(task.plan.as_slice()).and_then(|proto| {
+ proto.try_into_physical_plan(
+ executor.ctx.as_ref(),
+ codec.physical_extension_codec(),
+ )
+ })?;
+
let shuffle_output_partitioning =
parse_protobuf_hash_partitioning(task.output_partitioning.as_ref())?;
diff --git a/ballista/rust/executor/src/executor.rs b/ballista/rust/executor/src/executor.rs
index 6bf1aeb..5902897 100644
--- a/ballista/rust/executor/src/executor.rs
+++ b/ballista/rust/executor/src/executor.rs
@@ -27,7 +27,7 @@ use datafusion::error::DataFusionError;
use datafusion::execution::runtime_env::RuntimeEnv;
use datafusion::physical_plan::display::DisplayableExecutionPlan;
use datafusion::physical_plan::{ExecutionPlan, Partitioning};
-use datafusion::prelude::ExecutionConfig;
+use datafusion::prelude::{ExecutionConfig, ExecutionContext};
/// Ballista executor
pub struct Executor {
@@ -36,24 +36,30 @@ pub struct Executor {
/// Specification like total task slots
pub specification: ExecutorSpecification,
+
+ /// DataFusion execution context
+ pub ctx: Arc<ExecutionContext>,
}
impl Executor {
/// Create a new executor instance
- pub fn new(work_dir: &str) -> Self {
+ pub fn new(work_dir: &str, ctx: Arc<ExecutionContext>) -> Self {
Executor::new_with_specification(
work_dir,
ExecutorSpecification { task_slots: 4 },
+ ctx,
)
}
pub fn new_with_specification(
work_dir: &str,
specification: ExecutorSpecification,
+ ctx: Arc<ExecutionContext>,
) -> Self {
Self {
work_dir: work_dir.to_owned(),
specification,
+ ctx,
}
}
}
diff --git a/ballista/rust/executor/src/executor_server.rs b/ballista/rust/executor/src/executor_server.rs
index 3f220ea..7e4c9a1 100644
--- a/ballista/rust/executor/src/executor_server.rs
+++ b/ballista/rust/executor/src/executor_server.rs
@@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.
-use std::convert::TryInto;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::mpsc;
@@ -37,15 +36,17 @@ use ballista_core::serde::protobuf::{
UpdateTaskStatusParams,
};
use ballista_core::serde::scheduler::{ExecutorSpecification, ExecutorState};
+use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec};
use datafusion::physical_plan::ExecutionPlan;
use crate::as_task_status;
use crate::executor::Executor;
-pub async fn startup(
+pub async fn startup<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
mut scheduler: SchedulerGrpcClient<Channel>,
executor: Arc<Executor>,
executor_meta: ExecutorRegistration,
+ codec: BallistaCodec<T, U>,
) {
// TODO make the buffer size configurable
let (tx_task, rx_task) = mpsc::channel::<TaskDefinition>(1000);
@@ -55,6 +56,7 @@ pub async fn startup(
executor.clone(),
executor_meta.clone(),
ExecutorEnv { tx_task },
+ codec,
);
// 1. Start executor grpc service
@@ -126,12 +128,13 @@ async fn register_executor(
}
#[derive(Clone)]
-pub struct ExecutorServer {
+pub struct ExecutorServer<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> {
_start_time: u128,
executor: Arc<Executor>,
executor_meta: ExecutorRegistration,
scheduler: SchedulerGrpcClient<Channel>,
executor_env: ExecutorEnv,
+ codec: BallistaCodec<T, U>,
}
#[derive(Clone)]
@@ -141,12 +144,13 @@ struct ExecutorEnv {
unsafe impl Sync for ExecutorEnv {}
-impl ExecutorServer {
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T, U> {
fn new(
scheduler: SchedulerGrpcClient<Channel>,
executor: Arc<Executor>,
executor_meta: ExecutorRegistration,
executor_env: ExecutorEnv,
+ codec: BallistaCodec<T, U>,
) -> Self {
Self {
_start_time: SystemTime::now()
@@ -157,6 +161,7 @@ impl ExecutorServer {
executor_meta,
scheduler,
executor_env,
+ codec,
}
}
@@ -180,7 +185,16 @@ impl ExecutorServer {
);
info!("Start to run task {}", task_id_log);
- let plan: Arc<dyn ExecutionPlan> = (&task.plan.unwrap()).try_into().unwrap();
+ let encoded_plan = &task.plan.as_slice();
+
+ let plan: Arc<dyn ExecutionPlan> =
+ U::try_decode(encoded_plan).and_then(|proto| {
+ proto.try_into_physical_plan(
+ self.executor.ctx.as_ref(),
+ self.codec.physical_extension_codec(),
+ )
+ })?;
+
let shuffle_output_partitioning =
parse_protobuf_hash_partitioning(task.output_partitioning.as_ref())?;
@@ -221,12 +235,12 @@ impl ExecutorServer {
}
}
-struct Heartbeater {
- executor_server: Arc<ExecutorServer>,
+struct Heartbeater<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> {
+ executor_server: Arc<ExecutorServer<T, U>>,
}
-impl Heartbeater {
- fn new(executor_server: Arc<ExecutorServer>) -> Self {
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> Heartbeater<T, U> {
+ fn new(executor_server: Arc<ExecutorServer<T, U>>) -> Self {
Self { executor_server }
}
@@ -242,12 +256,12 @@ impl Heartbeater {
}
}
-struct TaskRunnerPool {
- executor_server: Arc<ExecutorServer>,
+struct TaskRunnerPool<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> {
+ executor_server: Arc<ExecutorServer<T, U>>,
}
-impl TaskRunnerPool {
- fn new(executor_server: Arc<ExecutorServer>) -> Self {
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskRunnerPool<T, U> {
+ fn new(executor_server: Arc<ExecutorServer<T, U>>) -> Self {
Self { executor_server }
}
@@ -269,7 +283,9 @@ impl TaskRunnerPool {
}
#[tonic::async_trait]
-impl ExecutorGrpc for ExecutorServer {
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorGrpc
+ for ExecutorServer<T, U>
+{
async fn launch_task(
&self,
request: Request<LaunchTaskParams>,
diff --git a/ballista/rust/executor/src/main.rs b/ballista/rust/executor/src/main.rs
index 2321ce3..3440f82 100644
--- a/ballista/rust/executor/src/main.rs
+++ b/ballista/rust/executor/src/main.rs
@@ -30,13 +30,15 @@ use uuid::Uuid;
use ballista_core::config::TaskSchedulingPolicy;
use ballista_core::serde::protobuf::{
executor_registration, scheduler_grpc_client::SchedulerGrpcClient,
- ExecutorRegistration,
+ ExecutorRegistration, LogicalPlanNode, PhysicalPlanNode,
};
use ballista_core::serde::scheduler::ExecutorSpecification;
+use ballista_core::serde::BallistaCodec;
use ballista_core::{print_version, BALLISTA_VERSION};
use ballista_executor::executor::Executor;
use ballista_executor::flight_service::BallistaFlightService;
use config::prelude::*;
+use datafusion::prelude::ExecutionContext;
#[macro_use]
extern crate configure_me;
@@ -102,15 +104,19 @@ async fn main() -> Result<()> {
let executor_specification = ExecutorSpecification {
task_slots: opt.concurrent_tasks as u32,
};
- let executor = Arc::new(Executor::new_with_specification(
+ let executor: Arc<Executor> = Arc::new(Executor::new_with_specification(
&work_dir,
executor_specification,
+ Arc::new(ExecutionContext::new()),
));
let scheduler = SchedulerGrpcClient::connect(scheduler_url)
.await
.context("Could not connect to scheduler")?;
+ let default_codec: BallistaCodec<LogicalPlanNode, PhysicalPlanNode> =
+ BallistaCodec::default();
+
let scheduler_policy = opt.task_scheduling_policy;
match scheduler_policy {
TaskSchedulingPolicy::PushStaged => {
@@ -118,6 +124,7 @@ async fn main() -> Result<()> {
scheduler,
executor.clone(),
executor_meta,
+ default_codec,
));
}
_ => {
@@ -126,6 +133,7 @@ async fn main() -> Result<()> {
executor.clone(),
executor_meta,
opt.concurrent_tasks,
+ default_codec,
));
}
}
diff --git a/ballista/rust/executor/src/standalone.rs b/ballista/rust/executor/src/standalone.rs
index 03f47bb..fa5f06e 100644
--- a/ballista/rust/executor/src/standalone.rs
+++ b/ballista/rust/executor/src/standalone.rs
@@ -18,12 +18,15 @@
use std::sync::Arc;
use arrow_flight::flight_service_server::FlightServiceServer;
+
+use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec};
use ballista_core::{
error::Result,
serde::protobuf::executor_registration::OptionalHost,
serde::protobuf::{scheduler_grpc_client::SchedulerGrpcClient, ExecutorRegistration},
BALLISTA_VERSION,
};
+use datafusion::prelude::ExecutionContext;
use log::info;
use tempfile::TempDir;
use tokio::net::TcpListener;
@@ -32,16 +35,21 @@ use uuid::Uuid;
use crate::{execution_loop, executor::Executor, flight_service::BallistaFlightService};
-pub async fn new_standalone_executor(
+pub async fn new_standalone_executor<
+ T: 'static + AsLogicalPlan,
+ U: 'static + AsExecutionPlan,
+>(
scheduler: SchedulerGrpcClient<Channel>,
concurrent_tasks: usize,
+ codec: BallistaCodec<T, U>,
) -> Result<()> {
let work_dir = TempDir::new()?
.into_path()
.into_os_string()
.into_string()
.unwrap();
- let executor = Arc::new(Executor::new(&work_dir));
+ let ctx = Arc::new(ExecutionContext::new());
+ let executor: Arc<Executor> = Arc::new(Executor::new(&work_dir, ctx));
let service = BallistaFlightService::new(executor.clone());
@@ -70,6 +78,7 @@ pub async fn new_standalone_executor(
executor,
executor_meta,
concurrent_tasks,
+ codec,
));
Ok(())
}
diff --git a/ballista/rust/scheduler/src/api/handlers.rs b/ballista/rust/scheduler/src/api/handlers.rs
index b65f4e2..f568221 100644
--- a/ballista/rust/scheduler/src/api/handlers.rs
+++ b/ballista/rust/scheduler/src/api/handlers.rs
@@ -11,6 +11,7 @@
// limitations under the License.
use crate::SchedulerServer;
+use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan};
use ballista_core::BALLISTA_VERSION;
use warp::Rejection;
@@ -29,8 +30,8 @@ pub struct ExecutorMetaResponse {
pub last_seen: u128,
}
-pub(crate) async fn scheduler_state(
- data_server: SchedulerServer,
+pub(crate) async fn scheduler_state<T: AsLogicalPlan, U: AsExecutionPlan>(
+ data_server: SchedulerServer<T, U>,
) -> Result<impl warp::Reply, Rejection> {
// TODO: Display last seen information in UI
let executors: Vec<ExecutorMetaResponse> = data_server
diff --git a/ballista/rust/scheduler/src/api/mod.rs b/ballista/rust/scheduler/src/api/mod.rs
index 45f281a..a0cd821 100644
--- a/ballista/rust/scheduler/src/api/mod.rs
+++ b/ballista/rust/scheduler/src/api/mod.rs
@@ -14,6 +14,7 @@ mod handlers;
use crate::SchedulerServer;
use anyhow::Result;
+use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan};
use std::{
pin::Pin,
task::{Context as TaskContext, Poll},
@@ -73,13 +74,16 @@ fn map_option_err<T, U: Into<Error>>(
err.map(|e| e.map_err(Into::into))
}
-fn with_data_server(
- db: SchedulerServer,
-) -> impl Filter<Extract = (SchedulerServer,), Error = std::convert::Infallible> + Clone {
+fn with_data_server<T: AsLogicalPlan + Clone, U: 'static + AsExecutionPlan>(
+ db: SchedulerServer<T, U>,
+) -> impl Filter<Extract = (SchedulerServer<T, U>,), Error = std::convert::Infallible> + Clone
+{
warp::any().map(move || db.clone())
}
-pub fn get_routes(scheduler_server: SchedulerServer) -> BoxedFilter<(impl Reply,)> {
+pub fn get_routes<T: AsLogicalPlan + Clone, U: 'static + AsExecutionPlan>(
+ scheduler_server: SchedulerServer<T, U>,
+) -> BoxedFilter<(impl Reply,)> {
let routes = warp::path("state")
.and(with_data_server(scheduler_server))
.and_then(handlers::scheduler_state);
diff --git a/ballista/rust/scheduler/src/lib.rs b/ballista/rust/scheduler/src/lib.rs
index cc508d3..3459cce 100644
--- a/ballista/rust/scheduler/src/lib.rs
+++ b/ballista/rust/scheduler/src/lib.rs
@@ -43,6 +43,7 @@ pub mod externalscaler {
use std::collections::{HashMap, HashSet};
use std::fmt;
+use std::marker::PhantomData;
use std::{convert::TryInto, sync::Arc};
use ballista_core::serde::protobuf::{
@@ -104,18 +105,21 @@ use ballista_core::error::BallistaError;
use ballista_core::execution_plans::ShuffleWriterExec;
use ballista_core::serde::protobuf::executor_grpc_client::ExecutorGrpcClient;
use ballista_core::serde::scheduler::to_proto::hash_partitioning_to_proto;
+use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec};
use datafusion::prelude::{ExecutionConfig, ExecutionContext};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::{mpsc, RwLock};
use tonic::transport::Channel;
#[derive(Clone)]
-pub struct SchedulerServer {
- pub(crate) state: Arc<SchedulerState>,
+pub struct SchedulerServer<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> {
+ pub(crate) state: Arc<SchedulerState<T, U>>,
start_time: u128,
policy: TaskSchedulingPolicy,
scheduler_env: Option<SchedulerEnv>,
executors_client: Arc<RwLock<HashMap<String, ExecutorGrpcClient<Channel>>>>,
+ ctx: Arc<RwLock<ExecutionContext>>,
+ codec: BallistaCodec<T, U>,
}
#[derive(Clone)]
@@ -123,13 +127,20 @@ pub struct SchedulerEnv {
pub tx_job: mpsc::Sender<String>,
}
-impl SchedulerServer {
- pub fn new(config: Arc<dyn ConfigBackendClient>, namespace: String) -> Self {
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T, U> {
+ pub fn new(
+ config: Arc<dyn ConfigBackendClient>,
+ namespace: String,
+ ctx: Arc<RwLock<ExecutionContext>>,
+ codec: BallistaCodec<T, U>,
+ ) -> Self {
SchedulerServer::new_with_policy(
config,
namespace,
TaskSchedulingPolicy::PullStaged,
None,
+ ctx,
+ codec,
)
}
@@ -138,8 +149,10 @@ impl SchedulerServer {
namespace: String,
policy: TaskSchedulingPolicy,
scheduler_env: Option<SchedulerEnv>,
+ ctx: Arc<RwLock<ExecutionContext>>,
+ codec: BallistaCodec<T, U>,
) -> Self {
- let state = Arc::new(SchedulerState::new(config, namespace));
+ let state = Arc::new(SchedulerState::new(config, namespace, codec.clone()));
let state_clone = state.clone();
// TODO: we should elect a leader in the scheduler cluster and run this only in the leader
@@ -154,6 +167,8 @@ impl SchedulerServer {
policy,
scheduler_env,
executors_client: Arc::new(RwLock::new(HashMap::new())),
+ ctx,
+ codec,
}
}
@@ -226,6 +241,7 @@ impl SchedulerServer {
ret.push(Vec::new());
}
let mut num_tasks = 0;
+ let ctx = self.ctx.read().await;
loop {
info!("Go inside fetching task loop");
let mut has_tasks = true;
@@ -235,7 +251,7 @@ impl SchedulerServer {
}
let plan = self
.state
- .assign_next_schedulable_job_task(&executor.executor_id, job_id)
+ .assign_next_schedulable_job_task(&executor.executor_id, job_id, &ctx)
.await
.map_err(|e| {
let msg = format!("Error finding next assignable task: {}", e);
@@ -266,8 +282,21 @@ impl SchedulerServer {
)));
};
+ let mut buf: Vec<u8> = vec![];
+ U::try_from_physical_plan(
+ plan,
+ self.codec.physical_extension_codec(),
+ )
+ .and_then(|m| m.try_encode(&mut buf))
+ .map_err(|e| {
+ Status::internal(format!(
+ "error serializing execution plan: {:?}",
+ e
+ ))
+ })?;
+
ret[idx].push(TaskDefinition {
- plan: Some(plan.try_into().unwrap()),
+ plan: buf,
task_id: status.partition_id,
output_partitioning: hash_partitioning_to_proto(
output_partitioning,
@@ -297,13 +326,19 @@ impl SchedulerServer {
}
}
-pub struct TaskScheduler {
- scheduler_server: Arc<SchedulerServer>,
+pub struct TaskScheduler<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> {
+ scheduler_server: Arc<SchedulerServer<T, U>>,
+ plan_repr: PhantomData<T>,
+ exec_repr: PhantomData<U>,
}
-impl TaskScheduler {
- pub fn new(scheduler_server: Arc<SchedulerServer>) -> Self {
- Self { scheduler_server }
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskScheduler<T, U> {
+ pub fn new(scheduler_server: Arc<SchedulerServer<T, U>>) -> Self {
+ Self {
+ scheduler_server,
+ plan_repr: PhantomData,
+ exec_repr: PhantomData,
+ }
}
pub fn start(&self, mut rx_job: mpsc::Receiver<String>) {
@@ -324,7 +359,9 @@ impl TaskScheduler {
const INFLIGHT_TASKS_METRIC_NAME: &str = "inflight_tasks";
#[tonic::async_trait]
-impl ExternalScaler for SchedulerServer {
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExternalScaler
+ for SchedulerServer<T, U>
+{
async fn is_active(
&self,
_request: Request<ScaledObjectRef>,
@@ -371,7 +408,9 @@ impl ExternalScaler for SchedulerServer {
}
#[tonic::async_trait]
-impl SchedulerGrpc for SchedulerServer {
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
+ for SchedulerServer<T, U>
+{
async fn poll_work(
&self,
request: Request<PollWorkParams>,
@@ -425,9 +464,10 @@ impl SchedulerGrpc for SchedulerServer {
})?;
}
let task: Result<Option<_>, Status> = if can_accept_task {
+ let ctx = self.ctx.read().await;
let plan = self
.state
- .assign_next_schedulable_task(&metadata.id)
+ .assign_next_schedulable_task(&metadata.id, &ctx)
.await
.map_err(|e| {
let msg = format!("Error finding next assignable task: {}", e);
@@ -457,8 +497,20 @@ impl SchedulerGrpc for SchedulerServer {
plan_clone
)));
};
+ let mut buf: Vec<u8> = vec![];
+ U::try_from_physical_plan(
+ plan,
+ self.codec.physical_extension_codec(),
+ )
+ .and_then(|m| m.try_encode(&mut buf))
+ .map_err(|e| {
+ Status::internal(format!(
+ "error serializing execution plan: {:?}",
+ e
+ ))
+ })?;
Ok(Some(TaskDefinition {
- plan: Some(plan.try_into().unwrap()),
+ plan: buf,
task_id: status.partition_id,
output_partitioning: hash_partitioning_to_proto(
output_partitioning,
@@ -722,33 +774,28 @@ impl SchedulerGrpc for SchedulerServer {
) -> std::result::Result<Response<ExecuteQueryResult>, tonic::Status> {
if let ExecuteQueryParams {
query: Some(query),
- settings,
+ settings: _,
} = request.into_inner()
{
- // parse config
- let mut config_builder = BallistaConfig::builder();
- for kv_pair in &settings {
- config_builder = config_builder.set(&kv_pair.key, &kv_pair.value);
- }
- let config = config_builder.build().map_err(|e| {
- let msg = format!("Could not parse configs: {}", e);
- error!("{}", msg);
- tonic::Status::internal(msg)
- })?;
-
let plan = match query {
- Query::LogicalPlan(logical_plan) => {
- // parse protobuf
- (&logical_plan).try_into().map_err(|e| {
- let msg = format!("Could not parse logical plan protobuf: {}", e);
- error!("{}", msg);
- tonic::Status::internal(msg)
- })?
+ Query::LogicalPlan(message) => {
+ let ctx = self.ctx.read().await;
+ T::try_decode(message.as_slice())
+ .and_then(|m| {
+ m.try_into_logical_plan(
+ &ctx,
+ self.codec.logical_extension_codec(),
+ )
+ })
+ .map_err(|e| {
+ let msg =
+ format!("Could not parse logical plan protobuf: {}", e);
+ error!("{}", msg);
+ tonic::Status::internal(msg)
+ })?
}
Query::Sql(sql) => {
- //TODO we can't just create a new context because we need a context that has
- // tables registered from previous SQL statements that have been executed
- let mut ctx = create_datafusion_context(&config);
+ let mut ctx = self.ctx.write().await;
let df = ctx.sql(&sql).await.map_err(|e| {
let msg = format!("Error parsing SQL: {}", e);
error!("{}", msg);
@@ -788,9 +835,9 @@ impl SchedulerGrpc for SchedulerServer {
Some(self.scheduler_env.as_ref().unwrap().tx_job.clone())
}
};
+ let datafusion_ctx = self.ctx.read().await.clone();
tokio::spawn(async move {
// create physical plan using DataFusion
- let datafusion_ctx = create_datafusion_context(&config);
macro_rules! fail_job {
($code :expr) => {{
match $code {
@@ -941,13 +988,17 @@ pub fn create_datafusion_context(config: &BallistaConfig) -> ExecutionContext {
#[cfg(all(test, feature = "sled"))]
mod test {
use std::sync::Arc;
+ use tokio::sync::RwLock;
use tonic::Request;
use ballista_core::error::BallistaError;
use ballista_core::serde::protobuf::{
- executor_registration::OptionalHost, ExecutorRegistration, PollWorkParams,
+ executor_registration::OptionalHost, ExecutorRegistration, LogicalPlanNode,
+ PhysicalPlanNode, PollWorkParams,
};
+ use ballista_core::serde::BallistaCodec;
+ use datafusion::prelude::ExecutionContext;
use super::{
state::{SchedulerState, StandaloneClient},
@@ -958,8 +1009,15 @@ mod test {
async fn test_poll_work() -> Result<(), BallistaError> {
let state = Arc::new(StandaloneClient::try_new_temporary()?);
let namespace = "default";
- let scheduler = SchedulerServer::new(state.clone(), namespace.to_owned());
- let state = SchedulerState::new(state, namespace.to_string());
+ let scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerServer::new(
+ state.clone(),
+ namespace.to_owned(),
+ Arc::new(RwLock::new(ExecutionContext::new())),
+ BallistaCodec::default(),
+ );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(state, namespace.to_string(), BallistaCodec::default());
let exec_meta = ExecutorRegistration {
id: "abc".to_owned(),
optional_host: Some(OptionalHost::Host("".to_owned())),
diff --git a/ballista/rust/scheduler/src/main.rs b/ballista/rust/scheduler/src/main.rs
index 5da5bbe..c906602 100644
--- a/ballista/rust/scheduler/src/main.rs
+++ b/ballista/rust/scheduler/src/main.rs
@@ -29,7 +29,10 @@ use tower::Service;
use ballista_core::BALLISTA_VERSION;
use ballista_core::{
- print_version, serde::protobuf::scheduler_grpc_server::SchedulerGrpcServer,
+ print_version,
+ serde::protobuf::{
+ scheduler_grpc_server::SchedulerGrpcServer, LogicalPlanNode, PhysicalPlanNode,
+ },
};
use ballista_scheduler::api::{get_routes, EitherBody, Error};
#[cfg(feature = "etcd")]
@@ -42,8 +45,9 @@ use ballista_scheduler::{
};
use ballista_core::config::TaskSchedulingPolicy;
+use ballista_core::serde::BallistaCodec;
use log::info;
-use tokio::sync::mpsc;
+use tokio::sync::{mpsc, RwLock};
#[macro_use]
extern crate configure_me;
@@ -59,6 +63,7 @@ mod config {
}
use config::prelude::*;
+use datafusion::prelude::ExecutionContext;
async fn start_server(
config_backend: Arc<dyn ConfigBackendClient>,
@@ -75,22 +80,31 @@ async fn start_server(
"Starting Scheduler grpc server with task scheduling policy of {:?}",
policy
);
- let scheduler_server = match policy {
- TaskSchedulingPolicy::PushStaged => {
- // TODO make the buffer size configurable
- let (tx_job, rx_job) = mpsc::channel::<String>(10000);
- let scheduler_server = SchedulerServer::new_with_policy(
+ let scheduler_server: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
+ match policy {
+ TaskSchedulingPolicy::PushStaged => {
+ // TODO make the buffer size configurable
+ let (tx_job, rx_job) = mpsc::channel::<String>(10000);
+ let scheduler_server = SchedulerServer::new_with_policy(
+ config_backend.clone(),
+ namespace.clone(),
+ policy,
+ Some(SchedulerEnv { tx_job }),
+ Arc::new(RwLock::new(ExecutionContext::new())),
+ BallistaCodec::default(),
+ );
+ let task_scheduler =
+ TaskScheduler::new(Arc::new(scheduler_server.clone()));
+ task_scheduler.start(rx_job);
+ scheduler_server
+ }
+ _ => SchedulerServer::new(
config_backend.clone(),
namespace.clone(),
- policy,
- Some(SchedulerEnv { tx_job }),
- );
- let task_scheduler = TaskScheduler::new(Arc::new(scheduler_server.clone()));
- task_scheduler.start(rx_job);
- scheduler_server
- }
- _ => SchedulerServer::new(config_backend.clone(), namespace.clone()),
- };
+ Arc::new(RwLock::new(ExecutionContext::new())),
+ BallistaCodec::default(),
+ ),
+ };
Ok(Server::bind(&addr)
.serve(make_service_fn(move |request: &AddrStream| {
diff --git a/ballista/rust/scheduler/src/planner.rs b/ballista/rust/scheduler/src/planner.rs
index fdd1435..68f26e9 100644
--- a/ballista/rust/scheduler/src/planner.rs
+++ b/ballista/rust/scheduler/src/planner.rs
@@ -250,7 +250,7 @@ mod test {
use crate::test_utils::datafusion_test_context;
use ballista_core::error::BallistaError;
use ballista_core::execution_plans::UnresolvedShuffleExec;
- use ballista_core::serde::protobuf;
+ use ballista_core::serde::{protobuf, AsExecutionPlan, BallistaCodec};
use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
use datafusion::physical_plan::hash_aggregate::{AggregateMode, HashAggregateExec};
use datafusion::physical_plan::hash_join::HashJoinExec;
@@ -259,7 +259,9 @@ mod test {
coalesce_partitions::CoalescePartitionsExec, projection::ProjectionExec,
};
use datafusion::physical_plan::{displayable, ExecutionPlan};
- use std::convert::TryInto;
+ use datafusion::prelude::ExecutionContext;
+
+ use ballista_core::serde::protobuf::{LogicalPlanNode, PhysicalPlanNode};
use std::sync::Arc;
use uuid::Uuid;
@@ -572,8 +574,16 @@ order by
fn roundtrip_operator(
plan: Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>, BallistaError> {
- let proto: protobuf::PhysicalPlanNode = plan.clone().try_into()?;
- let result_exec_plan: Arc<dyn ExecutionPlan> = (&proto).try_into()?;
+ let ctx = ExecutionContext::new();
+ let codec: BallistaCodec<LogicalPlanNode, PhysicalPlanNode> =
+ BallistaCodec::default();
+ let proto: protobuf::PhysicalPlanNode =
+ protobuf::PhysicalPlanNode::try_from_physical_plan(
+ plan.clone(),
+ codec.physical_extension_codec(),
+ )?;
+ let result_exec_plan: Arc<dyn ExecutionPlan> =
+ (&proto).try_into_physical_plan(&ctx, codec.physical_extension_codec())?;
Ok(result_exec_plan)
}
}
diff --git a/ballista/rust/scheduler/src/standalone.rs b/ballista/rust/scheduler/src/standalone.rs
index 55239d8..2d7250e 100644
--- a/ballista/rust/scheduler/src/standalone.rs
+++ b/ballista/rust/scheduler/src/standalone.rs
@@ -15,13 +15,19 @@
// specific language governing permissions and limitations
// under the License.
+use ballista_core::serde::BallistaCodec;
use ballista_core::{
- error::Result, serde::protobuf::scheduler_grpc_server::SchedulerGrpcServer,
+ error::Result,
+ serde::protobuf::{
+ scheduler_grpc_server::SchedulerGrpcServer, LogicalPlanNode, PhysicalPlanNode,
+ },
BALLISTA_VERSION,
};
+use datafusion::prelude::ExecutionContext;
use log::info;
use std::{net::SocketAddr, sync::Arc};
use tokio::net::TcpListener;
+use tokio::sync::RwLock;
use tonic::transport::Server;
use crate::{state::StandaloneClient, SchedulerServer};
@@ -29,10 +35,14 @@ use crate::{state::StandaloneClient, SchedulerServer};
pub async fn new_standalone_scheduler() -> Result<SocketAddr> {
let client = StandaloneClient::try_new_temporary()?;
- let server = SchedulerGrpcServer::new(SchedulerServer::new(
+ let server: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> = SchedulerServer::new(
Arc::new(client),
"ballista".to_string(),
- ));
+ Arc::new(RwLock::new(ExecutionContext::new())),
+ BallistaCodec::default(),
+ );
+
+ let server = SchedulerGrpcServer::new(server);
// Let the OS assign a random, free port
let listener = TcpListener::bind("localhost:0").await?;
let addr = listener.local_addr()?;
diff --git a/ballista/rust/scheduler/src/state/mod.rs b/ballista/rust/scheduler/src/state/mod.rs
index 45d915a..974516d 100644
--- a/ballista/rust/scheduler/src/state/mod.rs
+++ b/ballista/rust/scheduler/src/state/mod.rs
@@ -16,9 +16,7 @@
// under the License.
use std::time::{SystemTime, UNIX_EPOCH};
-use std::{
- any::type_name, collections::HashMap, convert::TryInto, sync::Arc, time::Duration,
-};
+use std::{any::type_name, collections::HashMap, sync::Arc, time::Duration};
use datafusion::physical_plan::ExecutionPlan;
use futures::{Stream, StreamExt};
@@ -28,12 +26,14 @@ use tokio::sync::OwnedMutexGuard;
use ballista_core::serde::protobuf::{
self, job_status, task_status, CompletedJob, CompletedTask, ExecutorHeartbeat,
- ExecutorMetadata, FailedJob, FailedTask, JobStatus, PhysicalPlanNode, RunningJob,
- RunningTask, TaskStatus,
+ ExecutorMetadata, FailedJob, FailedTask, JobStatus, RunningJob, RunningTask,
+ TaskStatus,
};
use ballista_core::serde::scheduler::{ExecutorData, PartitionStats};
+use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec};
use ballista_core::{error::BallistaError, serde::scheduler::ExecutorMeta};
use ballista_core::{error::Result, execution_plans::UnresolvedShuffleExec};
+use datafusion::prelude::ExecutionContext;
use super::planner::remove_unresolved_shuffles;
@@ -83,16 +83,23 @@ pub enum WatchEvent {
}
#[derive(Clone)]
-pub(super) struct SchedulerState {
+pub(super) struct SchedulerState<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
+{
config_client: Arc<dyn ConfigBackendClient>,
namespace: String,
+ codec: BallistaCodec<T, U>,
}
-impl SchedulerState {
- pub fn new(config_client: Arc<dyn ConfigBackendClient>, namespace: String) -> Self {
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T, U> {
+ pub fn new(
+ config_client: Arc<dyn ConfigBackendClient>,
+ namespace: String,
+ codec: BallistaCodec<T, U>,
+ ) -> Self {
Self {
config_client,
namespace,
+ codec,
}
}
@@ -271,8 +278,12 @@ impl SchedulerState {
) -> Result<()> {
let key = get_stage_plan_key(&self.namespace, job_id, stage_id);
let value = {
- let proto: PhysicalPlanNode = plan.try_into()?;
- encode_protobuf(&proto)?
+ let mut buf: Vec<u8> = vec![];
+ let proto =
+ U::try_from_physical_plan(plan, self.codec.physical_extension_codec())?;
+ proto.try_encode(&mut buf)?;
+
+ buf
};
self.config_client.clone().put(key, value).await
}
@@ -281,6 +292,7 @@ impl SchedulerState {
&self,
job_id: &str,
stage_id: usize,
+ ctx: &ExecutionContext,
) -> Result<Arc<dyn ExecutionPlan>> {
let key = get_stage_plan_key(&self.namespace, job_id, stage_id);
let value = &self.config_client.get(&key).await?;
@@ -290,8 +302,11 @@ impl SchedulerState {
key
)));
}
- let value: PhysicalPlanNode = decode_protobuf(value)?;
- Ok((&value).try_into()?)
+ let value = U::try_decode(value.as_slice())?;
+ let plan =
+ value.try_into_physical_plan(ctx, self.codec.physical_extension_codec())?;
+
+ Ok(plan)
}
pub async fn get_job_tasks(
@@ -352,9 +367,10 @@ impl SchedulerState {
pub async fn assign_next_schedulable_task(
&self,
executor_id: &str,
+ ctx: &ExecutionContext,
) -> Result<Option<(TaskStatus, Arc<dyn ExecutionPlan>)>> {
let tasks = self.get_all_tasks().await?;
- self.assign_next_schedulable_task_inner(executor_id, tasks)
+ self.assign_next_schedulable_task_inner(executor_id, tasks, ctx)
.await
}
@@ -362,9 +378,10 @@ impl SchedulerState {
&self,
executor_id: &str,
job_id: &str,
+ ctx: &ExecutionContext,
) -> Result<Option<(TaskStatus, Arc<dyn ExecutionPlan>)>> {
let job_tasks = self.get_job_tasks(job_id).await?;
- self.assign_next_schedulable_task_inner(executor_id, job_tasks)
+ self.assign_next_schedulable_task_inner(executor_id, job_tasks, ctx)
.await
}
@@ -372,8 +389,9 @@ impl SchedulerState {
&self,
executor_id: &str,
tasks: HashMap<String, TaskStatus>,
+ ctx: &ExecutionContext,
) -> Result<Option<(TaskStatus, Arc<dyn ExecutionPlan>)>> {
- match self.get_next_schedulable_task(tasks).await? {
+ match self.get_next_schedulable_task(tasks, ctx).await? {
Some((status, plan)) => {
let mut status = status.clone();
status.status = Some(task_status::Status::Running(RunningTask {
@@ -389,6 +407,7 @@ impl SchedulerState {
async fn get_next_schedulable_task(
&self,
tasks: HashMap<String, TaskStatus>,
+ ctx: &ExecutionContext,
) -> Result<Option<(TaskStatus, Arc<dyn ExecutionPlan>)>> {
// TODO: Make the duration a configurable parameter
let executors = self
@@ -398,7 +417,7 @@ impl SchedulerState {
if status.status.is_none() {
let partition = status.partition_id.as_ref().unwrap();
let plan = self
- .get_stage_plan(&partition.job_id, partition.stage_id as usize)
+ .get_stage_plan(&partition.job_id, partition.stage_id as usize, ctx)
.await?;
// Let's try to resolve any unresolved shuffles we find
@@ -778,9 +797,10 @@ mod test {
use std::sync::Arc;
use ballista_core::serde::protobuf::{
- job_status, task_status, CompletedTask, FailedTask, JobStatus, PartitionId,
- QueuedJob, RunningJob, RunningTask, TaskStatus,
+ job_status, task_status, CompletedTask, FailedTask, JobStatus, LogicalPlanNode,
+ PartitionId, PhysicalPlanNode, QueuedJob, RunningJob, RunningTask, TaskStatus,
};
+ use ballista_core::serde::BallistaCodec;
use ballista_core::{error::BallistaError, serde::scheduler::ExecutorMeta};
use super::{
@@ -790,10 +810,12 @@ mod test {
#[tokio::test]
async fn executor_metadata() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let meta = ExecutorMeta {
id: "123".to_owned(),
host: "localhost".to_owned(),
@@ -813,10 +835,12 @@ mod test {
#[tokio::test]
async fn job_metadata() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let meta = JobStatus {
status: Some(job_status::Status::Queued(QueuedJob {})),
};
@@ -832,10 +856,12 @@ mod test {
#[tokio::test]
async fn job_metadata_non_existant() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let meta = JobStatus {
status: Some(job_status::Status::Queued(QueuedJob {})),
};
@@ -847,10 +873,12 @@ mod test {
#[tokio::test]
async fn task_status() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let meta = TaskStatus {
status: Some(task_status::Status::Failed(FailedTask {
error: "error".to_owned(),
@@ -873,10 +901,12 @@ mod test {
#[tokio::test]
async fn task_status_non_existant() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let meta = TaskStatus {
status: Some(task_status::Status::Failed(FailedTask {
error: "error".to_owned(),
@@ -895,10 +925,12 @@ mod test {
#[tokio::test]
async fn task_synchronize_job_status_queued() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let job_id = "job";
let job_status = JobStatus {
status: Some(job_status::Status::Queued(QueuedJob {})),
@@ -912,10 +944,12 @@ mod test {
#[tokio::test]
async fn task_synchronize_job_status_running() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let job_id = "job";
let job_status = JobStatus {
status: Some(job_status::Status::Running(RunningJob {})),
@@ -952,10 +986,12 @@ mod test {
#[tokio::test]
async fn task_synchronize_job_status_running2() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let job_id = "job";
let job_status = JobStatus {
status: Some(job_status::Status::Running(RunningJob {})),
@@ -990,10 +1026,12 @@ mod test {
#[tokio::test]
async fn task_synchronize_job_status_completed() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let job_id = "job";
let job_status = JobStatus {
status: Some(job_status::Status::Running(RunningJob {})),
@@ -1034,10 +1072,12 @@ mod test {
#[tokio::test]
async fn task_synchronize_job_status_completed2() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let job_id = "job";
let job_status = JobStatus {
status: Some(job_status::Status::Queued(QueuedJob {})),
@@ -1078,10 +1118,12 @@ mod test {
#[tokio::test]
async fn task_synchronize_job_status_failed() -> Result<(), BallistaError> {
- let state = SchedulerState::new(
- Arc::new(StandaloneClient::try_new_temporary()?),
- "test".to_string(),
- );
+ let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerState::new(
+ Arc::new(StandaloneClient::try_new_temporary()?),
+ "test".to_string(),
+ BallistaCodec::default(),
+ );
let job_id = "job";
let job_status = JobStatus {
status: Some(job_status::Status::Running(RunningJob {})),
diff --git a/benchmarks/src/bin/tpch.rs b/benchmarks/src/bin/tpch.rs
index df34507..3880ca6 100644
--- a/benchmarks/src/bin/tpch.rs
+++ b/benchmarks/src/bin/tpch.rs
@@ -1361,15 +1361,20 @@ mod tests {
mod ballista_round_trip {
use super::*;
- use ballista_core::serde::protobuf;
+ use ballista_core::serde::{
+ protobuf, AsExecutionPlan, AsLogicalPlan, BallistaCodec,
+ };
use datafusion::physical_plan::ExecutionPlan;
- use std::convert::TryInto;
async fn round_trip_query(n: usize) -> Result<()> {
let config = ExecutionConfig::new()
.with_target_partitions(1)
.with_batch_size(10);
let mut ctx = ExecutionContext::with_config(config);
+ let codec: BallistaCodec<
+ protobuf::LogicalPlanNode,
+ protobuf::PhysicalPlanNode,
+ > = BallistaCodec::default();
// set tpch_data_path to dummy value and skip physical plan serde test when TPCH_DATA
// is not set.
@@ -1395,8 +1400,15 @@ mod tests {
// test logical plan round trip
let plan = create_logical_plan(&mut ctx, n)?;
- let proto: protobuf::LogicalPlanNode = (&plan).try_into().unwrap();
- let round_trip: LogicalPlan = (&proto).try_into().unwrap();
+ let proto: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ &plan,
+ codec.logical_extension_codec(),
+ )
+ .unwrap();
+ let round_trip: LogicalPlan = (&proto)
+ .try_into_logical_plan(&ctx, codec.logical_extension_codec())
+ .unwrap();
assert_eq!(
format!("{:?}", plan),
format!("{:?}", round_trip),
@@ -1405,8 +1417,15 @@ mod tests {
// test optimized logical plan round trip
let plan = ctx.optimize(&plan)?;
- let proto: protobuf::LogicalPlanNode = (&plan).try_into().unwrap();
- let round_trip: LogicalPlan = (&proto).try_into().unwrap();
+ let proto: protobuf::LogicalPlanNode =
+ protobuf::LogicalPlanNode::try_from_logical_plan(
+ &plan,
+ codec.logical_extension_codec(),
+ )
+ .unwrap();
+ let round_trip: LogicalPlan = (&proto)
+ .try_into_logical_plan(&ctx, codec.logical_extension_codec())
+ .unwrap();
assert_eq!(
format!("{:?}", plan),
format!("{:?}", round_trip),
@@ -1417,8 +1436,14 @@ mod tests {
if env::var("TPCH_DATA").is_ok() {
let physical_plan = ctx.create_physical_plan(&plan).await?;
let proto: protobuf::PhysicalPlanNode =
- (physical_plan.clone()).try_into().unwrap();
- let round_trip: Arc<dyn ExecutionPlan> = (&proto).try_into().unwrap();
+ protobuf::PhysicalPlanNode::try_from_physical_plan(
+ physical_plan.clone(),
+ codec.physical_extension_codec(),
+ )
+ .unwrap();
+ let round_trip: Arc<dyn ExecutionPlan> = (&proto)
+ .try_into_physical_plan(&ctx, codec.physical_extension_codec())
+ .unwrap();
assert_eq!(
format!("{:?}", physical_plan),
format!("{:?}", round_trip),