From c616f29ed79257336da4344156558104176c2fe2 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 21 Jul 2022 20:07:36 -0600 Subject: [PATCH 01/11] Can connect with Java FlightSQL client and hit get_flight_info_statement() breakpoint --- ballista/rust/scheduler/Cargo.toml | 1 + ballista/rust/scheduler/src/flight_sql.rs | 195 ++++++++++++++++++++++ ballista/rust/scheduler/src/lib.rs | 1 + ballista/rust/scheduler/src/main.rs | 5 + 4 files changed, 202 insertions(+) create mode 100644 ballista/rust/scheduler/src/flight_sql.rs diff --git a/ballista/rust/scheduler/Cargo.toml b/ballista/rust/scheduler/Cargo.toml index 8ff4093084..2587240e99 100644 --- a/ballista/rust/scheduler/Cargo.toml +++ b/ballista/rust/scheduler/Cargo.toml @@ -36,6 +36,7 @@ sled = ["sled_package", "tokio-stream"] [dependencies] anyhow = "1" +arrow-flight = { version = "18.0.0", features = ["flight-sql-experimental"] } async-recursion = "1.0.0" async-trait = "0.1.41" ballista-core = { path = "../core", version = "0.7.0" } diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs new file mode 100644 index 0000000000..be021ba53e --- /dev/null +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -0,0 +1,195 @@ +use arrow_flight::{FlightData, FlightDescriptor, FlightInfo}; +use arrow_flight::flight_service_server::FlightService; +use arrow_flight::sql::{ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTables, CommandGetTableTypes, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate, SqlInfo, TicketStatementQuery}; +use arrow_flight::sql::server::FlightSqlService; +use tonic::{Response, Status, Streaming}; + +#[derive(Clone)] +pub struct FlightSqlServiceImpl {} + +#[tonic::async_trait] +impl FlightSqlService for FlightSqlServiceImpl { + type FlightService = FlightSqlServiceImpl; + // get_flight_info + async fn get_flight_info_statement( + &self, + _query: CommandStatementQuery, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(tonic::Status::unimplemented("Not yet implemented")) + } + async fn get_flight_info_prepared_statement( + &self, + _query: CommandPreparedStatementQuery, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn get_flight_info_catalogs( + &self, + _query: CommandGetCatalogs, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn get_flight_info_schemas( + &self, + _query: CommandGetDbSchemas, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn get_flight_info_tables( + &self, + _query: CommandGetTables, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn get_flight_info_table_types( + &self, + _query: CommandGetTableTypes, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn get_flight_info_sql_info( + &self, + _query: CommandGetSqlInfo, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn get_flight_info_primary_keys( + &self, + _query: CommandGetPrimaryKeys, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn get_flight_info_exported_keys( + &self, + _query: CommandGetExportedKeys, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn get_flight_info_imported_keys( + &self, + _query: CommandGetImportedKeys, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn get_flight_info_cross_reference( + &self, + _query: CommandGetCrossReference, + _request: FlightDescriptor, + ) -> Result, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + // do_get + async fn do_get_statement( + &self, + _ticket: TicketStatementQuery, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + + async fn do_get_prepared_statement( + &self, + _query: CommandPreparedStatementQuery, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_get_catalogs( + &self, + _query: CommandGetCatalogs, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_get_schemas( + &self, + _query: CommandGetDbSchemas, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_get_tables( + &self, + _query: CommandGetTables, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_get_table_types( + &self, + _query: CommandGetTableTypes, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_get_sql_info( + &self, + _query: CommandGetSqlInfo, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_get_primary_keys( + &self, + _query: CommandGetPrimaryKeys, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_get_exported_keys( + &self, + _query: CommandGetExportedKeys, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_get_imported_keys( + &self, + _query: CommandGetImportedKeys, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_get_cross_reference( + &self, + _query: CommandGetCrossReference, + ) -> Result::DoGetStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + // do_put + async fn do_put_statement_update( + &self, + _ticket: CommandStatementUpdate, + ) -> Result { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_put_prepared_statement_query( + &self, + _query: CommandPreparedStatementQuery, + _request: Streaming, + ) -> Result::DoPutStream>, Status> { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_put_prepared_statement_update( + &self, + _query: CommandPreparedStatementUpdate, + _request: Streaming, + ) -> Result { + Err(Status::unimplemented("Not yet implemented")) + } + // do_action + async fn do_action_create_prepared_statement( + &self, + _query: ActionCreatePreparedStatementRequest, + ) -> Result { + Err(Status::unimplemented("Not yet implemented")) + } + async fn do_action_close_prepared_statement( + &self, + _query: ActionClosePreparedStatementRequest, + ) { + unimplemented!("Not yet implemented") + } + + async fn register_sql_info(&self, _id: i32, _result: &SqlInfo) {} +} diff --git a/ballista/rust/scheduler/src/lib.rs b/ballista/rust/scheduler/src/lib.rs index ea39ef02ef..e7bded18ed 100644 --- a/ballista/rust/scheduler/src/lib.rs +++ b/ballista/rust/scheduler/src/lib.rs @@ -26,3 +26,4 @@ pub mod state; #[cfg(test)] pub mod test_utils; +pub mod flight_sql; diff --git a/ballista/rust/scheduler/src/main.rs b/ballista/rust/scheduler/src/main.rs index 2dd34c92ee..4a951be573 100644 --- a/ballista/rust/scheduler/src/main.rs +++ b/ballista/rust/scheduler/src/main.rs @@ -23,6 +23,7 @@ use futures::future::{self, Either, TryFutureExt}; use hyper::{server::conn::AddrStream, service::make_service_fn, Server}; use std::convert::Infallible; use std::{net::SocketAddr, sync::Arc}; +use arrow_flight::flight_service_server::FlightServiceServer; use tonic::transport::server::Connected; use tonic::transport::Server as TonicServer; use tower::Service; @@ -61,6 +62,7 @@ mod config { use config::prelude::*; use datafusion::execution::context::default_session_builder; +use ballista_scheduler::flight_sql::FlightSqlServiceImpl; async fn start_server( config_backend: Arc, @@ -100,10 +102,13 @@ async fn start_server( let scheduler_grpc_server = SchedulerGrpcServer::new(scheduler_server.clone()); + let flight_sql_server = FlightServiceServer::new(FlightSqlServiceImpl {}); + let keda_scaler = ExternalScalerServer::new(scheduler_server.clone()); let mut tonic = TonicServer::builder() .add_service(scheduler_grpc_server) + .add_service(flight_sql_server) .add_service(keda_scaler) .into_service(); let mut warp = warp::service(get_routes(scheduler_server.clone())); From 7ebe0e89058c5a531b272b0f2bb522dcf371ecc3 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Fri, 22 Jul 2022 13:12:35 -0600 Subject: [PATCH 02/11] [WIP] trying to make a schema --- ballista/rust/scheduler/Cargo.toml | 1 + ballista/rust/scheduler/src/flight_sql.rs | 89 ++++++++++++++++++++++- ballista/rust/scheduler/src/main.rs | 4 +- 3 files changed, 89 insertions(+), 5 deletions(-) diff --git a/ballista/rust/scheduler/Cargo.toml b/ballista/rust/scheduler/Cargo.toml index 2587240e99..ca60f06174 100644 --- a/ballista/rust/scheduler/Cargo.toml +++ b/ballista/rust/scheduler/Cargo.toml @@ -46,6 +46,7 @@ datafusion = { git = "https://github.com/apache/arrow-datafusion", rev = "6a5de4 datafusion-proto = { git = "https://github.com/apache/arrow-datafusion", rev = "6a5de4fe08597896ab6375e3e4b76c5744dcfba7" } env_logger = "0.9" etcd-client = { version = "0.9", optional = true } +flatbuffers = { version = "2.1.2" } futures = "0.3" http = "0.2" http-body = "0.4" diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index be021ba53e..ea0c599a9d 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -1,11 +1,34 @@ -use arrow_flight::{FlightData, FlightDescriptor, FlightInfo}; +use arrow_flight::{FlightData, FlightDescriptor, FlightInfo, IpcMessage}; +use arrow_flight::flight_descriptor::DescriptorType; use arrow_flight::flight_service_server::FlightService; use arrow_flight::sql::{ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTables, CommandGetTableTypes, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate, SqlInfo, TicketStatementQuery}; use arrow_flight::sql::server::FlightSqlService; +use datafusion::arrow::array::StringBuilder; +use datafusion::arrow::ipc; +use datafusion::arrow::ipc::{FieldBuilder, SchemaBuilder}; +use datafusion::parquet::data_type::AsBytes; +use log::debug; use tonic::{Response, Status, Streaming}; +use prost::Message; + +use crate::scheduler_server::SchedulerServer; +use datafusion_proto::protobuf::LogicalPlanNode; +use flatbuffers::{FlatBufferBuilder, Vector}; +use ballista_core::{ + serde::protobuf::{PhysicalPlanNode}, +}; +use ballista_core::config::BallistaConfig; #[derive(Clone)] -pub struct FlightSqlServiceImpl {} +pub struct FlightSqlServiceImpl { + server: SchedulerServer, +} + +impl FlightSqlServiceImpl { + pub fn new(server: SchedulerServer) -> Self { + Self { server } + } +} #[tonic::async_trait] impl FlightSqlService for FlightSqlServiceImpl { @@ -13,10 +36,68 @@ impl FlightSqlService for FlightSqlServiceImpl { // get_flight_info async fn get_flight_info_statement( &self, - _query: CommandStatementQuery, + query: CommandStatementQuery, _request: FlightDescriptor, ) -> Result, Status> { - Err(tonic::Status::unimplemented("Not yet implemented")) + debug!("Got query:\n{}", query.query); + + // Run query + let config_builder = BallistaConfig::builder(); + let config = config_builder.build() + .map_err(|e| Status::internal(format!("Error building config: {}", e)))?; + let ctx = self.server + .state + .session_manager + .create_session(&config) + .await + .map_err(|e| { + Status::internal(format!( + "Failed to create SessionContext: {:?}", + e + )) + })?; + let plan = ctx + .sql(&query.query.as_str()) + .await + .and_then(|df| df.to_logical_plan()) + .map_err(|e| Status::internal(format!("Error building plan: {}", e)))?; + + let mut fbb = FlatBufferBuilder::new(); + + let mut sb = SchemaBuilder::new(&mut fbb); + + fbb.start_vector(0); + let mut fb = FieldBuilder::new(&mut fbb); + let field_name = fbb.create_string("EXPR$0"); + fb.add_name(field_name); + let field = fb.finish(); + + fbb.push() + + let fields = fbb.create_vector(&[field]); + + sb.add_fields(fields); + let schema = sb.finish(); + let bytes = Vec::from(schema.to_le_bytes()); + + // Generate response + let mut buf = vec![]; + let ticket = TicketStatementQuery { statement_handle: vec![] }; + ticket.encode(&mut buf) + .map_err(|e| Status::internal(format!("Error encoding message: {}", e)))?; + let total_records = 0; + let total_bytes = 0; + let ep = vec![]; + let flight_desc = FlightDescriptor { + r#type: DescriptorType::Cmd.into(), + cmd: bytes, + path: vec![] + }; + let msg = IpcMessage(buf); + let info = FlightInfo::new(msg, Some(flight_desc), ep, total_records, total_bytes); + let resp = Response::new(info); + debug!("Responding to query..."); + Ok(resp) } async fn get_flight_info_prepared_statement( &self, diff --git a/ballista/rust/scheduler/src/main.rs b/ballista/rust/scheduler/src/main.rs index 4a951be573..c617f1e53d 100644 --- a/ballista/rust/scheduler/src/main.rs +++ b/ballista/rust/scheduler/src/main.rs @@ -102,7 +102,9 @@ async fn start_server( let scheduler_grpc_server = SchedulerGrpcServer::new(scheduler_server.clone()); - let flight_sql_server = FlightServiceServer::new(FlightSqlServiceImpl {}); + let flight_sql_server = FlightServiceServer::new( + FlightSqlServiceImpl::new(scheduler_server.clone()) + ); let keda_scaler = ExternalScalerServer::new(scheduler_server.clone()); From 5e5850fae87ae18a81fa879fe5dd7581ad91ebfb Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Fri, 22 Jul 2022 14:23:13 -0600 Subject: [PATCH 03/11] Can generate a pretty good-looking response --- ballista/rust/scheduler/src/flight_sql.rs | 62 ++++++++++------------- 1 file changed, 28 insertions(+), 34 deletions(-) diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index ea0c599a9d..3768a37343 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -1,23 +1,21 @@ -use arrow_flight::{FlightData, FlightDescriptor, FlightInfo, IpcMessage}; +use arrow_flight::{FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, Location, Ticket}; use arrow_flight::flight_descriptor::DescriptorType; use arrow_flight::flight_service_server::FlightService; use arrow_flight::sql::{ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTables, CommandGetTableTypes, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate, SqlInfo, TicketStatementQuery}; use arrow_flight::sql::server::FlightSqlService; -use datafusion::arrow::array::StringBuilder; -use datafusion::arrow::ipc; -use datafusion::arrow::ipc::{FieldBuilder, SchemaBuilder}; -use datafusion::parquet::data_type::AsBytes; use log::debug; use tonic::{Response, Status, Streaming}; -use prost::Message; use crate::scheduler_server::SchedulerServer; use datafusion_proto::protobuf::LogicalPlanNode; -use flatbuffers::{FlatBufferBuilder, Vector}; use ballista_core::{ serde::protobuf::{PhysicalPlanNode}, }; use ballista_core::config::BallistaConfig; +use arrow_flight::SchemaAsIpc; +use datafusion::arrow; +use datafusion::arrow::datatypes::Schema; +use datafusion::arrow::ipc::writer::{IpcDataGenerator, IpcWriteOptions}; #[derive(Clone)] pub struct FlightSqlServiceImpl { @@ -37,7 +35,7 @@ impl FlightSqlService for FlightSqlServiceImpl { async fn get_flight_info_statement( &self, query: CommandStatementQuery, - _request: FlightDescriptor, + request: FlightDescriptor, ) -> Result, Status> { debug!("Got query:\n{}", query.query); @@ -62,39 +60,35 @@ impl FlightSqlService for FlightSqlServiceImpl { .and_then(|df| df.to_logical_plan()) .map_err(|e| Status::internal(format!("Error building plan: {}", e)))?; - let mut fbb = FlatBufferBuilder::new(); - - let mut sb = SchemaBuilder::new(&mut fbb); - - fbb.start_vector(0); - let mut fb = FieldBuilder::new(&mut fbb); - let field_name = fbb.create_string("EXPR$0"); - fb.add_name(field_name); - let field = fb.finish(); + // transform schema + let arrow_schema: Schema = (&**plan.schema()).into(); + let options = IpcWriteOptions::default(); + let pair = SchemaAsIpc::new(&arrow_schema, &options); - fbb.push() - - let fields = fbb.create_vector(&[field]); - - sb.add_fields(fields); - let schema = sb.finish(); - let bytes = Vec::from(schema.to_le_bytes()); + let data_gen = IpcDataGenerator::default(); + let encoded_data = data_gen.schema_to_bytes(pair.0, pair.1); + let mut schema_bytes = vec![]; + arrow::ipc::writer::write_message(&mut schema_bytes, encoded_data, pair.1) + .map_err(|e| Status::internal(format!("Error encoding schema: {}", e)))?; // Generate response - let mut buf = vec![]; - let ticket = TicketStatementQuery { statement_handle: vec![] }; - ticket.encode(&mut buf) - .map_err(|e| Status::internal(format!("Error encoding message: {}", e)))?; - let total_records = 0; - let total_bytes = 0; - let ep = vec![]; + let ticket = Ticket { ticket: vec![] }; + let fiep = FlightEndpoint { + ticket: Some(ticket), + location: vec![Location { uri: "grpc+tcp://0.0.0.0:50050".to_string() }], + }; let flight_desc = FlightDescriptor { r#type: DescriptorType::Cmd.into(), - cmd: bytes, + cmd: vec![], path: vec![] }; - let msg = IpcMessage(buf); - let info = FlightInfo::new(msg, Some(flight_desc), ep, total_records, total_bytes); + let info = FlightInfo { + schema: schema_bytes, + flight_descriptor: Some(flight_desc), + endpoint: vec![fiep], + total_records: -1, + total_bytes: -1, + }; let resp = Response::new(info); debug!("Responding to query..."); Ok(resp) From 4a17a3abdbe9fc06800bd8a7ecf20a658bf8c764 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Fri, 22 Jul 2022 15:13:38 -0600 Subject: [PATCH 04/11] Getting to do_get_statement() with statement_handle --- ballista/rust/scheduler/src/flight_sql.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index 3768a37343..e661a38ee1 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -1,10 +1,11 @@ use arrow_flight::{FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, Location, Ticket}; use arrow_flight::flight_descriptor::DescriptorType; use arrow_flight::flight_service_server::FlightService; -use arrow_flight::sql::{ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTables, CommandGetTableTypes, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate, SqlInfo, TicketStatementQuery}; +use arrow_flight::sql::{ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTables, CommandGetTableTypes, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate, ProstAnyExt, ProstMessageExt, SqlInfo, TicketStatementQuery}; use arrow_flight::sql::server::FlightSqlService; use log::debug; use tonic::{Response, Status, Streaming}; +use prost::Message; use crate::scheduler_server::SchedulerServer; use datafusion_proto::protobuf::LogicalPlanNode; @@ -72,7 +73,11 @@ impl FlightSqlService for FlightSqlServiceImpl { .map_err(|e| Status::internal(format!("Error encoding schema: {}", e)))?; // Generate response - let ticket = Ticket { ticket: vec![] }; + let ticket = TicketStatementQuery { statement_handle: vec![1,2,3] }; + // let as_any = ProstAnyExt::pack(&ticket) + // .map_err(|e| Status::internal(format!("Error encoding schema: {}", e)))?; + let buf = ticket.as_any().encode_to_vec(); + let ticket = Ticket { ticket: buf }; let fiep = FlightEndpoint { ticket: Some(ticket), location: vec![Location { uri: "grpc+tcp://0.0.0.0:50050".to_string() }], From eb78d65f8fd6b5852a9f06ce034e6e7ecda71459 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Fri, 22 Jul 2022 15:36:39 -0600 Subject: [PATCH 05/11] Can retrieve plan --- ballista/rust/scheduler/Cargo.toml | 2 +- ballista/rust/scheduler/src/flight_sql.rs | 24 +++++++++++++++++------ 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/ballista/rust/scheduler/Cargo.toml b/ballista/rust/scheduler/Cargo.toml index ca60f06174..e0555a4624 100644 --- a/ballista/rust/scheduler/Cargo.toml +++ b/ballista/rust/scheduler/Cargo.toml @@ -63,11 +63,11 @@ tokio = { version = "1.0", features = ["full"] } tokio-stream = { version = "0.1", features = ["net"], optional = true } tonic = "0.7" tower = { version = "0.4" } +uuid = { version = "1.0", features = ["v4"] } warp = "0.3" [dev-dependencies] ballista-core = { path = "../core", version = "0.7.0" } -uuid = { version = "1.0", features = ["v4"] } [build-dependencies] configure_me_codegen = "0.4.1" diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index e661a38ee1..76b13129bb 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -1,7 +1,9 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; use arrow_flight::{FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, Location, Ticket}; use arrow_flight::flight_descriptor::DescriptorType; use arrow_flight::flight_service_server::FlightService; -use arrow_flight::sql::{ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTables, CommandGetTableTypes, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate, ProstAnyExt, ProstMessageExt, SqlInfo, TicketStatementQuery}; +use arrow_flight::sql::{ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTables, CommandGetTableTypes, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate, ProstMessageExt, SqlInfo, TicketStatementQuery}; use arrow_flight::sql::server::FlightSqlService; use log::debug; use tonic::{Response, Status, Streaming}; @@ -17,15 +19,18 @@ use arrow_flight::SchemaAsIpc; use datafusion::arrow; use datafusion::arrow::datatypes::Schema; use datafusion::arrow::ipc::writer::{IpcDataGenerator, IpcWriteOptions}; +use datafusion::logical_expr::LogicalPlan; +use uuid::{Uuid}; #[derive(Clone)] pub struct FlightSqlServiceImpl { server: SchedulerServer, + statements: Arc>>, } impl FlightSqlServiceImpl { pub fn new(server: SchedulerServer) -> Self { - Self { server } + Self { server, statements: Arc::new(Mutex::new(HashMap::new())) } } } @@ -73,9 +78,11 @@ impl FlightSqlService for FlightSqlServiceImpl { .map_err(|e| Status::internal(format!("Error encoding schema: {}", e)))?; // Generate response - let ticket = TicketStatementQuery { statement_handle: vec![1,2,3] }; - // let as_any = ProstAnyExt::pack(&ticket) - // .map_err(|e| Status::internal(format!("Error encoding schema: {}", e)))?; + let handle = Uuid::new_v4(); + let ticket = TicketStatementQuery { statement_handle: handle.as_bytes().to_vec() }; + let mut statements = self.statements.try_lock() + .map_err(|e| Status::internal(format!("Error locking statements: {}", e)))?; + statements.insert(handle, plan); let buf = ticket.as_any().encode_to_vec(); let ticket = Ticket { ticket: buf }; let fiep = FlightEndpoint { @@ -171,8 +178,13 @@ impl FlightSqlService for FlightSqlServiceImpl { // do_get async fn do_get_statement( &self, - _ticket: TicketStatementQuery, + ticket: TicketStatementQuery, ) -> Result::DoGetStream>, Status> { + let handle = Uuid::from_slice(&ticket.statement_handle) + .map_err(|e| Status::internal(format!("Error decoding ticket: {}", e)))?; + let statements = self.statements.try_lock() + .map_err(|e| Status::internal(format!("Error decoding ticket: {}", e)))?; + let plan = statements.get(&handle); Err(Status::unimplemented("Not yet implemented")) } From 1f4f5f54fb9b6895bfa639b453b914447f3484fe Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Sun, 24 Jul 2022 13:29:42 -0600 Subject: [PATCH 06/11] Can get results --- ballista/rust/scheduler/src/flight_sql.rs | 159 ++++++++++++++++++---- 1 file changed, 136 insertions(+), 23 deletions(-) diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index 76b13129bb..f4bb3dc501 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -1,13 +1,13 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex}; +use std::time::Duration; use arrow_flight::{FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, Location, Ticket}; use arrow_flight::flight_descriptor::DescriptorType; use arrow_flight::flight_service_server::FlightService; -use arrow_flight::sql::{ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTables, CommandGetTableTypes, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate, ProstMessageExt, SqlInfo, TicketStatementQuery}; +use arrow_flight::sql::{ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTables, CommandGetTableTypes, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate, SqlInfo, TicketStatementQuery}; use arrow_flight::sql::server::FlightSqlService; -use log::debug; +use log::{debug, error}; use tonic::{Response, Status, Streaming}; -use prost::Message; use crate::scheduler_server::SchedulerServer; use datafusion_proto::protobuf::LogicalPlanNode; @@ -20,7 +20,13 @@ use datafusion::arrow; use datafusion::arrow::datatypes::Schema; use datafusion::arrow::ipc::writer::{IpcDataGenerator, IpcWriteOptions}; use datafusion::logical_expr::LogicalPlan; +use tokio::time::sleep; use uuid::{Uuid}; +use crate::scheduler_server::event::QueryStageSchedulerEvent; +use ballista_core::serde::protobuf::job_status; +use ballista_core::serde::protobuf::JobStatus; +use ballista_core::serde::protobuf; +use prost::Message; #[derive(Clone)] pub struct FlightSqlServiceImpl { @@ -66,6 +72,124 @@ impl FlightSqlService for FlightSqlServiceImpl { .and_then(|df| df.to_logical_plan()) .map_err(|e| Status::internal(format!("Error building plan: {}", e)))?; + // enqueue job + let job_id = self.server.state.task_manager.generate_job_id(); + + self.server.state + .task_manager + .queue_job(&job_id) + .await + .map_err(|e| { + let msg = format!("Failed to queue job {}: {:?}", job_id, e); + error!("{}", msg); + + Status::internal(msg) + })?; + + let query_stage_event_sender = + self.server.query_stage_event_loop.get_sender().map_err(|e| { + Status::internal(format!( + "Could not get query stage event sender due to: {}", + e + )) + })?; + + query_stage_event_sender + .post_event(QueryStageSchedulerEvent::JobQueued { + job_id: job_id.clone(), + session_id: ctx.session_id().clone(), + session_ctx: ctx, + plan: Box::new(plan.clone()), + }) + .await + .map_err(|e| { + let msg = + format!("Failed to send JobQueued event for {}: {:?}", job_id, e); + error!("{}", msg); + Status::internal(msg) + })?; + + // let handle = Uuid::new_v4(); + // let mut statements = self.statements.try_lock() + // .map_err(|e| Status::internal(format!("Error locking statements: {}", e)))?; + // statements.insert(handle, plan.clone()); + + // poll for job completion + let mut num_rows = 0; + let mut num_bytes = 0; + let fieps = loop { + sleep(Duration::from_millis(100)).await; + let status = self.server.state.task_manager.get_job_status(&job_id).await + .map_err(|e| { + let msg = format!("Error getting status for job {}: {:?}", job_id, e); + error!("{}", msg); + Status::internal(msg) + })?; + let status: JobStatus = match status { + Some(status) => status, + None => { + let msg = format!("Error getting status for job {}!", job_id); + error!("{}", msg); + Err(Status::internal(msg))? + } + }; + let status: job_status::Status = match status.status { + Some(status) => status, + None => { + let msg = format!("Error getting status for job {}!", job_id); + error!("{}", msg); + Err(Status::internal(msg))? + } + }; + let completed = match status { + job_status::Status::Queued(_) => continue, + job_status::Status::Running(_) => continue, + job_status::Status::Failed(e) => { + Err(Status::internal(format!("Error building plan: {}", e.error)))? + } + job_status::Status::Completed(comp) => comp + }; + let mut fieps: Vec<_> = vec![]; + for loc in completed.partition_location.iter() { + let fetch = if let Some(ref id) = loc.partition_id { + let fetch = protobuf::FetchPartition { + job_id: id.job_id.clone(), + stage_id: id.stage_id, + partition_id: id.partition_id, + path: loc.path.clone() + }; + protobuf::Action { + action_type: Some(protobuf::action::ActionType::FetchPartition(fetch)), + settings: vec![], + } + } else { + Err(Status::internal(format!("Error getting partition it")))? + }; + let authority = if let Some(ref md) = loc.executor_meta { + // pub id: ::prost::alloc::string::String, + format!("{}:{}", md.host, md.port) + } else { + Err(Status::internal(format!("Error getting location")))? + }; + if let Some(ref stats) = loc.partition_stats { + num_rows += stats.num_rows; + num_bytes += stats.num_bytes; + // pub num_batches: i64, + } else { + Err(Status::internal(format!("Error getting stats")))? + } + let loc = Location { uri: format!("grpc+tcp://{}", authority) }; + let buf = fetch.encode_to_vec(); + let ticket = Ticket { ticket: buf }; + let fiep = FlightEndpoint { + ticket: Some(ticket), + location: vec![loc], + }; + fieps.push(fiep); + } + break fieps; + }; + // transform schema let arrow_schema: Schema = (&**plan.schema()).into(); let options = IpcWriteOptions::default(); @@ -78,28 +202,17 @@ impl FlightSqlService for FlightSqlServiceImpl { .map_err(|e| Status::internal(format!("Error encoding schema: {}", e)))?; // Generate response - let handle = Uuid::new_v4(); - let ticket = TicketStatementQuery { statement_handle: handle.as_bytes().to_vec() }; - let mut statements = self.statements.try_lock() - .map_err(|e| Status::internal(format!("Error locking statements: {}", e)))?; - statements.insert(handle, plan); - let buf = ticket.as_any().encode_to_vec(); - let ticket = Ticket { ticket: buf }; - let fiep = FlightEndpoint { - ticket: Some(ticket), - location: vec![Location { uri: "grpc+tcp://0.0.0.0:50050".to_string() }], - }; let flight_desc = FlightDescriptor { r#type: DescriptorType::Cmd.into(), cmd: vec![], - path: vec![] + path: vec![], }; let info = FlightInfo { schema: schema_bytes, flight_descriptor: Some(flight_desc), - endpoint: vec![fiep], - total_records: -1, - total_bytes: -1, + endpoint: fieps, + total_records: num_rows, + total_bytes: num_bytes, }; let resp = Response::new(info); debug!("Responding to query..."); @@ -180,11 +293,11 @@ impl FlightSqlService for FlightSqlServiceImpl { &self, ticket: TicketStatementQuery, ) -> Result::DoGetStream>, Status> { - let handle = Uuid::from_slice(&ticket.statement_handle) - .map_err(|e| Status::internal(format!("Error decoding ticket: {}", e)))?; - let statements = self.statements.try_lock() - .map_err(|e| Status::internal(format!("Error decoding ticket: {}", e)))?; - let plan = statements.get(&handle); + // let handle = Uuid::from_slice(&ticket.statement_handle) + // .map_err(|e| Status::internal(format!("Error decoding ticket: {}", e)))?; + // let statements = self.statements.try_lock() + // .map_err(|e| Status::internal(format!("Error decoding ticket: {}", e)))?; + // let plan = statements.get(&handle); Err(Status::unimplemented("Not yet implemented")) } From 09ca27fe7c5bcf4c8261785d228fc7e47145dd46 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Sun, 24 Jul 2022 17:36:09 -0600 Subject: [PATCH 07/11] Helpful error messages --- ballista/rust/scheduler/src/flight_sql.rs | 54 ++++++++++++----------- 1 file changed, 28 insertions(+), 26 deletions(-) diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index f4bb3dc501..e11c11c52c 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -223,70 +223,71 @@ impl FlightSqlService for FlightSqlServiceImpl { _query: CommandPreparedStatementQuery, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement get_flight_info_prepared_statement")) } async fn get_flight_info_catalogs( &self, _query: CommandGetCatalogs, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement get_flight_info_catalogs")) } async fn get_flight_info_schemas( &self, _query: CommandGetDbSchemas, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement get_flight_info_schemas")) } async fn get_flight_info_tables( &self, _query: CommandGetTables, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement get_flight_info_tables")) } async fn get_flight_info_table_types( &self, _query: CommandGetTableTypes, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement get_flight_info_table_types")) } async fn get_flight_info_sql_info( &self, _query: CommandGetSqlInfo, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Not yet implemented")) + // TODO: implement for FlightSQL JDBC to work + Err(Status::unimplemented("Implement CommandGetSqlInfo")) } async fn get_flight_info_primary_keys( &self, _query: CommandGetPrimaryKeys, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement get_flight_info_primary_keys")) } async fn get_flight_info_exported_keys( &self, _query: CommandGetExportedKeys, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement get_flight_info_exported_keys")) } async fn get_flight_info_imported_keys( &self, _query: CommandGetImportedKeys, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement get_flight_info_imported_keys")) } async fn get_flight_info_cross_reference( &self, _query: CommandGetCrossReference, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement get_flight_info_cross_reference")) } // do_get async fn do_get_statement( @@ -298,102 +299,103 @@ impl FlightSqlService for FlightSqlServiceImpl { // let statements = self.statements.try_lock() // .map_err(|e| Status::internal(format!("Error decoding ticket: {}", e)))?; // let plan = statements.get(&handle); - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_statement")) } async fn do_get_prepared_statement( &self, _query: CommandPreparedStatementQuery, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_prepared_statement")) } async fn do_get_catalogs( &self, _query: CommandGetCatalogs, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_catalogs")) } async fn do_get_schemas( &self, _query: CommandGetDbSchemas, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_schemas")) } async fn do_get_tables( &self, _query: CommandGetTables, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_tables")) } async fn do_get_table_types( &self, _query: CommandGetTableTypes, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_table_types")) } async fn do_get_sql_info( &self, _query: CommandGetSqlInfo, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_sql_info")) } async fn do_get_primary_keys( &self, _query: CommandGetPrimaryKeys, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_primary_keys")) } async fn do_get_exported_keys( &self, _query: CommandGetExportedKeys, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_exported_keys")) } async fn do_get_imported_keys( &self, _query: CommandGetImportedKeys, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_imported_keys")) } async fn do_get_cross_reference( &self, _query: CommandGetCrossReference, ) -> Result::DoGetStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_get_cross_reference")) } // do_put async fn do_put_statement_update( &self, _ticket: CommandStatementUpdate, ) -> Result { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_put_statement_update")) } async fn do_put_prepared_statement_query( &self, _query: CommandPreparedStatementQuery, _request: Streaming, ) -> Result::DoPutStream>, Status> { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_put_prepared_statement_query")) } async fn do_put_prepared_statement_update( &self, _query: CommandPreparedStatementUpdate, _request: Streaming, ) -> Result { - Err(Status::unimplemented("Not yet implemented")) + Err(Status::unimplemented("Implement do_put_prepared_statement_update")) } // do_action async fn do_action_create_prepared_statement( &self, _query: ActionCreatePreparedStatementRequest, ) -> Result { - Err(Status::unimplemented("Not yet implemented")) + // TODO: implement for Flight SQL JDBC driver to work + Err(Status::unimplemented("Implement do_action_create_prepared_statement")) } async fn do_action_close_prepared_statement( &self, _query: ActionClosePreparedStatementRequest, ) { - unimplemented!("Not yet implemented") + unimplemented!("Implement do_action_close_prepared_statement") } async fn register_sql_info(&self, _id: i32, _result: &SqlInfo) {} From 50c63c90ea478093e5f2dbe79a5de60ae62f0e57 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 25 Jul 2022 09:55:12 -0600 Subject: [PATCH 08/11] PR feedback --- ballista/rust/scheduler/src/flight_sql.rs | 164 +++++++++++++++------- ballista/rust/scheduler/src/lib.rs | 2 +- ballista/rust/scheduler/src/main.rs | 10 +- 3 files changed, 117 insertions(+), 59 deletions(-) diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index e11c11c52c..f50c4a06e3 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -1,42 +1,68 @@ -use std::collections::HashMap; -use std::sync::{Arc, Mutex}; -use std::time::Duration; -use arrow_flight::{FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, Location, Ticket}; +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + use arrow_flight::flight_descriptor::DescriptorType; use arrow_flight::flight_service_server::FlightService; -use arrow_flight::sql::{ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTables, CommandGetTableTypes, CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate, SqlInfo, TicketStatementQuery}; use arrow_flight::sql::server::FlightSqlService; +use arrow_flight::sql::{ + ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest, + ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference, + CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, + CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, + CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery, + CommandStatementUpdate, SqlInfo, TicketStatementQuery, +}; +use arrow_flight::{ + FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, Location, Ticket, +}; use log::{debug, error}; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use std::time::Duration; use tonic::{Response, Status, Streaming}; +use crate::scheduler_server::event::QueryStageSchedulerEvent; use crate::scheduler_server::SchedulerServer; -use datafusion_proto::protobuf::LogicalPlanNode; -use ballista_core::{ - serde::protobuf::{PhysicalPlanNode}, -}; -use ballista_core::config::BallistaConfig; use arrow_flight::SchemaAsIpc; +use ballista_core::config::BallistaConfig; +use ballista_core::serde::protobuf; +use ballista_core::serde::protobuf::job_status; +use ballista_core::serde::protobuf::JobStatus; +use ballista_core::serde::protobuf::PhysicalPlanNode; use datafusion::arrow; use datafusion::arrow::datatypes::Schema; use datafusion::arrow::ipc::writer::{IpcDataGenerator, IpcWriteOptions}; use datafusion::logical_expr::LogicalPlan; -use tokio::time::sleep; -use uuid::{Uuid}; -use crate::scheduler_server::event::QueryStageSchedulerEvent; -use ballista_core::serde::protobuf::job_status; -use ballista_core::serde::protobuf::JobStatus; -use ballista_core::serde::protobuf; +use datafusion_proto::protobuf::LogicalPlanNode; use prost::Message; +use tokio::time::sleep; +use uuid::Uuid; -#[derive(Clone)] pub struct FlightSqlServiceImpl { server: SchedulerServer, - statements: Arc>>, + _statements: Arc>>, } impl FlightSqlServiceImpl { pub fn new(server: SchedulerServer) -> Self { - Self { server, statements: Arc::new(Mutex::new(HashMap::new())) } + Self { + server, + _statements: Arc::new(Mutex::new(HashMap::new())), + } } } @@ -47,27 +73,26 @@ impl FlightSqlService for FlightSqlServiceImpl { async fn get_flight_info_statement( &self, query: CommandStatementQuery, - request: FlightDescriptor, + _request: FlightDescriptor, ) -> Result, Status> { debug!("Got query:\n{}", query.query); // Run query let config_builder = BallistaConfig::builder(); - let config = config_builder.build() + let config = config_builder + .build() .map_err(|e| Status::internal(format!("Error building config: {}", e)))?; - let ctx = self.server + let ctx = self + .server .state .session_manager .create_session(&config) .await .map_err(|e| { - Status::internal(format!( - "Failed to create SessionContext: {:?}", - e - )) + Status::internal(format!("Failed to create SessionContext: {:?}", e)) })?; let plan = ctx - .sql(&query.query.as_str()) + .sql(query.query.as_str()) .await .and_then(|df| df.to_logical_plan()) .map_err(|e| Status::internal(format!("Error building plan: {}", e)))?; @@ -75,7 +100,8 @@ impl FlightSqlService for FlightSqlServiceImpl { // enqueue job let job_id = self.server.state.task_manager.generate_job_id(); - self.server.state + self.server + .state .task_manager .queue_job(&job_id) .await @@ -86,8 +112,11 @@ impl FlightSqlService for FlightSqlServiceImpl { Status::internal(msg) })?; - let query_stage_event_sender = - self.server.query_stage_event_loop.get_sender().map_err(|e| { + let query_stage_event_sender = self + .server + .query_stage_event_loop + .get_sender() + .map_err(|e| { Status::internal(format!( "Could not get query stage event sender due to: {}", e @@ -119,7 +148,12 @@ impl FlightSqlService for FlightSqlServiceImpl { let mut num_bytes = 0; let fieps = loop { sleep(Duration::from_millis(100)).await; - let status = self.server.state.task_manager.get_job_status(&job_id).await + let status = self + .server + .state + .task_manager + .get_job_status(&job_id) + .await .map_err(|e| { let msg = format!("Error getting status for job {}: {:?}", job_id, e); error!("{}", msg); @@ -144,10 +178,11 @@ impl FlightSqlService for FlightSqlServiceImpl { let completed = match status { job_status::Status::Queued(_) => continue, job_status::Status::Running(_) => continue, - job_status::Status::Failed(e) => { - Err(Status::internal(format!("Error building plan: {}", e.error)))? - } - job_status::Status::Completed(comp) => comp + job_status::Status::Failed(e) => Err(Status::internal(format!( + "Error building plan: {}", + e.error + )))?, + job_status::Status::Completed(comp) => comp, }; let mut fieps: Vec<_> = vec![]; for loc in completed.partition_location.iter() { @@ -156,29 +191,34 @@ impl FlightSqlService for FlightSqlServiceImpl { job_id: id.job_id.clone(), stage_id: id.stage_id, partition_id: id.partition_id, - path: loc.path.clone() + path: loc.path.clone(), }; protobuf::Action { - action_type: Some(protobuf::action::ActionType::FetchPartition(fetch)), + action_type: Some(protobuf::action::ActionType::FetchPartition( + fetch, + )), settings: vec![], } } else { - Err(Status::internal(format!("Error getting partition it")))? + Err(Status::internal("Error getting partition ID".to_string()))? }; let authority = if let Some(ref md) = loc.executor_meta { - // pub id: ::prost::alloc::string::String, format!("{}:{}", md.host, md.port) } else { - Err(Status::internal(format!("Error getting location")))? + Err(Status::internal( + "Invalid partition location, missing executor metadata" + .to_string(), + ))? }; if let Some(ref stats) = loc.partition_stats { num_rows += stats.num_rows; num_bytes += stats.num_bytes; - // pub num_batches: i64, } else { - Err(Status::internal(format!("Error getting stats")))? + Err(Status::internal("Error getting stats".to_string()))? } - let loc = Location { uri: format!("grpc+tcp://{}", authority) }; + let loc = Location { + uri: format!("grpc+tcp://{}", authority), + }; let buf = fetch.encode_to_vec(); let ticket = Ticket { ticket: buf }; let fiep = FlightEndpoint { @@ -223,7 +263,9 @@ impl FlightSqlService for FlightSqlServiceImpl { _query: CommandPreparedStatementQuery, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Implement get_flight_info_prepared_statement")) + Err(Status::unimplemented( + "Implement get_flight_info_prepared_statement", + )) } async fn get_flight_info_catalogs( &self, @@ -251,7 +293,9 @@ impl FlightSqlService for FlightSqlServiceImpl { _query: CommandGetTableTypes, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Implement get_flight_info_table_types")) + Err(Status::unimplemented( + "Implement get_flight_info_table_types", + )) } async fn get_flight_info_sql_info( &self, @@ -266,33 +310,41 @@ impl FlightSqlService for FlightSqlServiceImpl { _query: CommandGetPrimaryKeys, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Implement get_flight_info_primary_keys")) + Err(Status::unimplemented( + "Implement get_flight_info_primary_keys", + )) } async fn get_flight_info_exported_keys( &self, _query: CommandGetExportedKeys, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Implement get_flight_info_exported_keys")) + Err(Status::unimplemented( + "Implement get_flight_info_exported_keys", + )) } async fn get_flight_info_imported_keys( &self, _query: CommandGetImportedKeys, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Implement get_flight_info_imported_keys")) + Err(Status::unimplemented( + "Implement get_flight_info_imported_keys", + )) } async fn get_flight_info_cross_reference( &self, _query: CommandGetCrossReference, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented("Implement get_flight_info_cross_reference")) + Err(Status::unimplemented( + "Implement get_flight_info_cross_reference", + )) } // do_get async fn do_get_statement( &self, - ticket: TicketStatementQuery, + _ticket: TicketStatementQuery, ) -> Result::DoGetStream>, Status> { // let handle = Uuid::from_slice(&ticket.statement_handle) // .map_err(|e| Status::internal(format!("Error decoding ticket: {}", e)))?; @@ -374,14 +426,18 @@ impl FlightSqlService for FlightSqlServiceImpl { _query: CommandPreparedStatementQuery, _request: Streaming, ) -> Result::DoPutStream>, Status> { - Err(Status::unimplemented("Implement do_put_prepared_statement_query")) + Err(Status::unimplemented( + "Implement do_put_prepared_statement_query", + )) } async fn do_put_prepared_statement_update( &self, _query: CommandPreparedStatementUpdate, _request: Streaming, ) -> Result { - Err(Status::unimplemented("Implement do_put_prepared_statement_update")) + Err(Status::unimplemented( + "Implement do_put_prepared_statement_update", + )) } // do_action async fn do_action_create_prepared_statement( @@ -389,7 +445,9 @@ impl FlightSqlService for FlightSqlServiceImpl { _query: ActionCreatePreparedStatementRequest, ) -> Result { // TODO: implement for Flight SQL JDBC driver to work - Err(Status::unimplemented("Implement do_action_create_prepared_statement")) + Err(Status::unimplemented( + "Implement do_action_create_prepared_statement", + )) } async fn do_action_close_prepared_statement( &self, diff --git a/ballista/rust/scheduler/src/lib.rs b/ballista/rust/scheduler/src/lib.rs index e7bded18ed..2c2fa4100a 100644 --- a/ballista/rust/scheduler/src/lib.rs +++ b/ballista/rust/scheduler/src/lib.rs @@ -24,6 +24,6 @@ pub mod scheduler_server; pub mod standalone; pub mod state; +pub mod flight_sql; #[cfg(test)] pub mod test_utils; -pub mod flight_sql; diff --git a/ballista/rust/scheduler/src/main.rs b/ballista/rust/scheduler/src/main.rs index c617f1e53d..a83353147b 100644 --- a/ballista/rust/scheduler/src/main.rs +++ b/ballista/rust/scheduler/src/main.rs @@ -18,12 +18,12 @@ //! Ballista Rust scheduler binary. use anyhow::{Context, Result}; +use arrow_flight::flight_service_server::FlightServiceServer; use ballista_scheduler::scheduler_server::externalscaler::external_scaler_server::ExternalScalerServer; use futures::future::{self, Either, TryFutureExt}; use hyper::{server::conn::AddrStream, service::make_service_fn, Server}; use std::convert::Infallible; use std::{net::SocketAddr, sync::Arc}; -use arrow_flight::flight_service_server::FlightServiceServer; use tonic::transport::server::Connected; use tonic::transport::Server as TonicServer; use tower::Service; @@ -60,9 +60,9 @@ mod config { )); } +use ballista_scheduler::flight_sql::FlightSqlServiceImpl; use config::prelude::*; use datafusion::execution::context::default_session_builder; -use ballista_scheduler::flight_sql::FlightSqlServiceImpl; async fn start_server( config_backend: Arc, @@ -102,9 +102,9 @@ async fn start_server( let scheduler_grpc_server = SchedulerGrpcServer::new(scheduler_server.clone()); - let flight_sql_server = FlightServiceServer::new( - FlightSqlServiceImpl::new(scheduler_server.clone()) - ); + let flight_sql_server = FlightServiceServer::new(FlightSqlServiceImpl::new( + scheduler_server.clone(), + )); let keda_scaler = ExternalScalerServer::new(scheduler_server.clone()); From bc30a9689586b4cbc696d20cf8145de752490a9b Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 25 Jul 2022 12:47:54 -0600 Subject: [PATCH 09/11] Refactor for sanity --- ballista/rust/scheduler/src/flight_sql.rs | 229 +++++++++++++--------- 1 file changed, 133 insertions(+), 96 deletions(-) diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index f50c4a06e3..beb7947692 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -41,12 +41,14 @@ use arrow_flight::SchemaAsIpc; use ballista_core::config::BallistaConfig; use ballista_core::serde::protobuf; use ballista_core::serde::protobuf::job_status; +use ballista_core::serde::protobuf::CompletedJob; use ballista_core::serde::protobuf::JobStatus; use ballista_core::serde::protobuf::PhysicalPlanNode; use datafusion::arrow; use datafusion::arrow::datatypes::Schema; use datafusion::arrow::ipc::writer::{IpcDataGenerator, IpcWriteOptions}; use datafusion::logical_expr::LogicalPlan; +use datafusion::prelude::SessionContext; use datafusion_proto::protobuf::LogicalPlanNode; use prost::Message; use tokio::time::sleep; @@ -64,20 +66,8 @@ impl FlightSqlServiceImpl { _statements: Arc::new(Mutex::new(HashMap::new())), } } -} -#[tonic::async_trait] -impl FlightSqlService for FlightSqlServiceImpl { - type FlightService = FlightSqlServiceImpl; - // get_flight_info - async fn get_flight_info_statement( - &self, - query: CommandStatementQuery, - _request: FlightDescriptor, - ) -> Result, Status> { - debug!("Got query:\n{}", query.query); - - // Run query + async fn create_ctx(&self) -> Result, Status> { let config_builder = BallistaConfig::builder(); let config = config_builder .build() @@ -91,15 +81,117 @@ impl FlightSqlService for FlightSqlServiceImpl { .map_err(|e| { Status::internal(format!("Failed to create SessionContext: {:?}", e)) })?; + Ok(ctx) + } + + async fn prepare_statement( + query: CommandStatementQuery, + ctx: &Arc, + ) -> Result { let plan = ctx .sql(query.query.as_str()) .await .and_then(|df| df.to_logical_plan()) .map_err(|e| Status::internal(format!("Error building plan: {}", e)))?; + Ok(plan) + } - // enqueue job - let job_id = self.server.state.task_manager.generate_job_id(); + async fn check_job(&self, job_id: &String) -> Result, Status> { + let status = self + .server + .state + .task_manager + .get_job_status(job_id) + .await + .map_err(|e| { + let msg = format!("Error getting status for job {}: {:?}", job_id, e); + error!("{}", msg); + Status::internal(msg) + })?; + let status: JobStatus = match status { + Some(status) => status, + None => { + let msg = format!("Error getting status for job {}!", job_id); + error!("{}", msg); + Err(Status::internal(msg))? + } + }; + let status: job_status::Status = match status.status { + Some(status) => status, + None => { + let msg = format!("Error getting status for job {}!", job_id); + error!("{}", msg); + Err(Status::internal(msg))? + } + }; + match status { + job_status::Status::Queued(_) => Ok(None), + job_status::Status::Running(_) => Ok(None), + job_status::Status::Failed(e) => Err(Status::internal(format!( + "Error building plan: {}", + e.error + )))?, + job_status::Status::Completed(comp) => Ok(Some(comp)), + } + } + + async fn job_to_fetch_part( + &self, + completed: CompletedJob, + num_rows: &mut i64, + num_bytes: &mut i64, + ) -> Result, Status> { + let mut fieps: Vec<_> = vec![]; + for loc in completed.partition_location.iter() { + let fetch = if let Some(ref id) = loc.partition_id { + let fetch = protobuf::FetchPartition { + job_id: id.job_id.clone(), + stage_id: id.stage_id, + partition_id: id.partition_id, + path: loc.path.clone(), + }; + protobuf::Action { + action_type: Some(protobuf::action::ActionType::FetchPartition( + fetch, + )), + settings: vec![], + } + } else { + Err(Status::internal("Error getting partition ID".to_string()))? + }; + let authority = if let Some(ref md) = loc.executor_meta { + format!("{}:{}", md.host, md.port) + } else { + Err(Status::internal( + "Invalid partition location, missing executor metadata".to_string(), + ))? + }; + if let Some(ref stats) = loc.partition_stats { + *num_rows += stats.num_rows; + *num_bytes += stats.num_bytes; + } else { + Err(Status::internal("Error getting stats".to_string()))? + } + let loc = Location { + uri: format!("grpc+tcp://{}", authority), + }; + let buf = fetch.encode_to_vec(); + let ticket = Ticket { ticket: buf }; + let fiep = FlightEndpoint { + ticket: Some(ticket), + location: vec![loc], + }; + fieps.push(fiep); + } + Ok(fieps) + } + async fn enqueue_job( + &self, + ctx: Arc, + plan: &LogicalPlan, + ) -> Result { + let job_id = self.server.state.task_manager.generate_job_id(); self.server .state .task_manager @@ -111,7 +203,6 @@ impl FlightSqlService for FlightSqlServiceImpl { Status::internal(msg) })?; - let query_stage_event_sender = self .server .query_stage_event_loop @@ -122,7 +213,6 @@ impl FlightSqlService for FlightSqlServiceImpl { e )) })?; - query_stage_event_sender .post_event(QueryStageSchedulerEvent::JobQueued { job_id: job_id.clone(), @@ -137,6 +227,24 @@ impl FlightSqlService for FlightSqlServiceImpl { error!("{}", msg); Status::internal(msg) })?; + Ok(job_id) + } +} + +#[tonic::async_trait] +impl FlightSqlService for FlightSqlServiceImpl { + type FlightService = FlightSqlServiceImpl; + // get_flight_info + async fn get_flight_info_statement( + &self, + query: CommandStatementQuery, + _request: FlightDescriptor, + ) -> Result, Status> { + debug!("Got query:\n{}", query.query); + + let ctx = self.create_ctx().await?; + let plan = Self::prepare_statement(query, &ctx).await?; + let job_id = self.enqueue_job(ctx, &plan).await?; // let handle = Uuid::new_v4(); // let mut statements = self.statements.try_lock() @@ -148,85 +256,14 @@ impl FlightSqlService for FlightSqlServiceImpl { let mut num_bytes = 0; let fieps = loop { sleep(Duration::from_millis(100)).await; - let status = self - .server - .state - .task_manager - .get_job_status(&job_id) - .await - .map_err(|e| { - let msg = format!("Error getting status for job {}: {:?}", job_id, e); - error!("{}", msg); - Status::internal(msg) - })?; - let status: JobStatus = match status { - Some(status) => status, - None => { - let msg = format!("Error getting status for job {}!", job_id); - error!("{}", msg); - Err(Status::internal(msg))? - } - }; - let status: job_status::Status = match status.status { - Some(status) => status, - None => { - let msg = format!("Error getting status for job {}!", job_id); - error!("{}", msg); - Err(Status::internal(msg))? - } + let completed = if let Some(comp) = self.check_job(&job_id).await? { + comp + } else { + continue; }; - let completed = match status { - job_status::Status::Queued(_) => continue, - job_status::Status::Running(_) => continue, - job_status::Status::Failed(e) => Err(Status::internal(format!( - "Error building plan: {}", - e.error - )))?, - job_status::Status::Completed(comp) => comp, - }; - let mut fieps: Vec<_> = vec![]; - for loc in completed.partition_location.iter() { - let fetch = if let Some(ref id) = loc.partition_id { - let fetch = protobuf::FetchPartition { - job_id: id.job_id.clone(), - stage_id: id.stage_id, - partition_id: id.partition_id, - path: loc.path.clone(), - }; - protobuf::Action { - action_type: Some(protobuf::action::ActionType::FetchPartition( - fetch, - )), - settings: vec![], - } - } else { - Err(Status::internal("Error getting partition ID".to_string()))? - }; - let authority = if let Some(ref md) = loc.executor_meta { - format!("{}:{}", md.host, md.port) - } else { - Err(Status::internal( - "Invalid partition location, missing executor metadata" - .to_string(), - ))? - }; - if let Some(ref stats) = loc.partition_stats { - num_rows += stats.num_rows; - num_bytes += stats.num_bytes; - } else { - Err(Status::internal("Error getting stats".to_string()))? - } - let loc = Location { - uri: format!("grpc+tcp://{}", authority), - }; - let buf = fetch.encode_to_vec(); - let ticket = Ticket { ticket: buf }; - let fiep = FlightEndpoint { - ticket: Some(ticket), - location: vec![loc], - }; - fieps.push(fiep); - } + let fieps = self + .job_to_fetch_part(completed, &mut num_rows, &mut num_bytes) + .await?; break fieps; }; @@ -234,7 +271,6 @@ impl FlightSqlService for FlightSqlServiceImpl { let arrow_schema: Schema = (&**plan.schema()).into(); let options = IpcWriteOptions::default(); let pair = SchemaAsIpc::new(&arrow_schema, &options); - let data_gen = IpcDataGenerator::default(); let encoded_data = data_gen.schema_to_bytes(pair.0, pair.1); let mut schema_bytes = vec![]; @@ -258,6 +294,7 @@ impl FlightSqlService for FlightSqlServiceImpl { debug!("Responding to query..."); Ok(resp) } + async fn get_flight_info_prepared_statement( &self, _query: CommandPreparedStatementQuery, From c9e35ecb884132f2d30f01673a73be6eb29e2b76 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 25 Jul 2022 13:51:55 -0600 Subject: [PATCH 10/11] Implement more of spec --- ballista/rust/scheduler/src/flight_sql.rs | 195 +++++++++++++++------- 1 file changed, 132 insertions(+), 63 deletions(-) diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index beb7947692..c7d8daa523 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -29,7 +29,7 @@ use arrow_flight::sql::{ use arrow_flight::{ FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, Location, Ticket, }; -use log::{debug, error}; +use log::{debug, error, warn}; use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -47,6 +47,7 @@ use ballista_core::serde::protobuf::PhysicalPlanNode; use datafusion::arrow; use datafusion::arrow::datatypes::Schema; use datafusion::arrow::ipc::writer::{IpcDataGenerator, IpcWriteOptions}; +use datafusion::common::DFSchemaRef; use datafusion::logical_expr::LogicalPlan; use datafusion::prelude::SessionContext; use datafusion_proto::protobuf::LogicalPlanNode; @@ -56,14 +57,14 @@ use uuid::Uuid; pub struct FlightSqlServiceImpl { server: SchedulerServer, - _statements: Arc>>, + statements: Arc>>, } impl FlightSqlServiceImpl { pub fn new(server: SchedulerServer) -> Self { Self { server, - _statements: Arc::new(Mutex::new(HashMap::new())), + statements: Arc::new(Mutex::new(HashMap::new())), } } @@ -85,11 +86,11 @@ impl FlightSqlServiceImpl { } async fn prepare_statement( - query: CommandStatementQuery, + query: &String, ctx: &Arc, ) -> Result { let plan = ctx - .sql(query.query.as_str()) + .sql(query.as_str()) .await .and_then(|df| df.to_logical_plan()) .map_err(|e| Status::internal(format!("Error building plan: {}", e)))?; @@ -127,11 +128,16 @@ impl FlightSqlServiceImpl { match status { job_status::Status::Queued(_) => Ok(None), job_status::Status::Running(_) => Ok(None), - job_status::Status::Failed(e) => Err(Status::internal(format!( - "Error building plan: {}", - e.error - )))?, - job_status::Status::Completed(comp) => Ok(Some(comp)), + job_status::Status::Failed(e) => { + warn!("Error executing plan: {:?}", e); + Err(Status::internal(format!( + "Error executing plan: {}", + e.error + )))? + }, + job_status::Status::Completed(comp) => { + Ok(Some(comp)) + }, } } @@ -186,6 +192,44 @@ impl FlightSqlServiceImpl { Ok(fieps) } + fn cache_plan(&self, plan: LogicalPlan) -> Result { + let handle = Uuid::new_v4(); + let mut statements = self.statements.try_lock() + .map_err(|e| Status::internal(format!("Error locking statements: {}", e)))?; + statements.insert(handle, plan); + Ok(handle) + } + + fn get_plan(&self, handle: &Uuid) -> Result { + let statements = self.statements.try_lock() + .map_err(|e| Status::internal(format!("Error locking statements: {}", e)))?; + let plan = if let Some(plan) = statements.get(&handle) { + plan + } else { + Err(Status::internal(format!("Statement handle not found: {}", handle)))? + }; + Ok(plan.clone()) + } + + fn remove_plan(&self, handle: Uuid) -> Result<(), Status> { + let mut statements = self.statements.try_lock() + .map_err(|e| Status::internal(format!("Error locking statements: {}", e)))?; + statements.remove(&handle); + Ok(()) + } + + fn df_schema_to_arrow(&self, schema: &DFSchemaRef) -> Result, Status> { + let arrow_schema: Schema = (&**schema).into(); + let options = IpcWriteOptions::default(); + let pair = SchemaAsIpc::new(&arrow_schema, &options); + let data_gen = IpcDataGenerator::default(); + let encoded_data = data_gen.schema_to_bytes(pair.0, pair.1); + let mut schema_bytes = vec![]; + arrow::ipc::writer::write_message(&mut schema_bytes, encoded_data, pair.1) + .map_err(|e| Status::internal(format!("Error encoding schema: {}", e)))?; + Ok(schema_bytes) + } + async fn enqueue_job( &self, ctx: Arc, @@ -229,28 +273,36 @@ impl FlightSqlServiceImpl { })?; Ok(job_id) } -} -#[tonic::async_trait] -impl FlightSqlService for FlightSqlServiceImpl { - type FlightService = FlightSqlServiceImpl; - // get_flight_info - async fn get_flight_info_statement( + fn create_resp( + schema_bytes: Vec, + fieps: Vec, + num_rows: i64, + num_bytes: i64, + ) -> Response { + let flight_desc = FlightDescriptor { + r#type: DescriptorType::Cmd.into(), + cmd: vec![], + path: vec![], + }; + let info = FlightInfo { + schema: schema_bytes, + flight_descriptor: Some(flight_desc), + endpoint: fieps, + total_records: num_rows, + total_bytes: num_bytes, + }; + let resp = Response::new(info); + resp + } + + async fn execute_plan( &self, - query: CommandStatementQuery, - _request: FlightDescriptor, + ctx: Arc, + plan: &LogicalPlan ) -> Result, Status> { - debug!("Got query:\n{}", query.query); - - let ctx = self.create_ctx().await?; - let plan = Self::prepare_statement(query, &ctx).await?; let job_id = self.enqueue_job(ctx, &plan).await?; - // let handle = Uuid::new_v4(); - // let mut statements = self.statements.try_lock() - // .map_err(|e| Status::internal(format!("Error locking statements: {}", e)))?; - // statements.insert(handle, plan.clone()); - // poll for job completion let mut num_rows = 0; let mut num_bytes = 0; @@ -267,43 +319,47 @@ impl FlightSqlService for FlightSqlServiceImpl { break fieps; }; - // transform schema - let arrow_schema: Schema = (&**plan.schema()).into(); - let options = IpcWriteOptions::default(); - let pair = SchemaAsIpc::new(&arrow_schema, &options); - let data_gen = IpcDataGenerator::default(); - let encoded_data = data_gen.schema_to_bytes(pair.0, pair.1); - let mut schema_bytes = vec![]; - arrow::ipc::writer::write_message(&mut schema_bytes, encoded_data, pair.1) - .map_err(|e| Status::internal(format!("Error encoding schema: {}", e)))?; - // Generate response - let flight_desc = FlightDescriptor { - r#type: DescriptorType::Cmd.into(), - cmd: vec![], - path: vec![], - }; - let info = FlightInfo { - schema: schema_bytes, - flight_descriptor: Some(flight_desc), - endpoint: fieps, - total_records: num_rows, - total_bytes: num_bytes, - }; - let resp = Response::new(info); + let schema_bytes = self.df_schema_to_arrow(plan.schema())?; + let resp = Self::create_resp(schema_bytes, fieps, num_rows, num_bytes); + Ok(resp) + } +} + +#[tonic::async_trait] +impl FlightSqlService for FlightSqlServiceImpl { + type FlightService = FlightSqlServiceImpl; + + async fn get_flight_info_statement( + &self, + query: CommandStatementQuery, + _request: FlightDescriptor, + ) -> Result, Status> { + debug!("Got query:\n{}", query.query); + + let ctx = self.create_ctx().await?; + let plan = Self::prepare_statement(&query.query, &ctx).await?; + let resp = self.execute_plan(ctx, &plan).await?; + debug!("Responding to query..."); Ok(resp) } async fn get_flight_info_prepared_statement( &self, - _query: CommandPreparedStatementQuery, + handle: CommandPreparedStatementQuery, _request: FlightDescriptor, ) -> Result, Status> { - Err(Status::unimplemented( - "Implement get_flight_info_prepared_statement", - )) + let ctx = self.create_ctx().await?; + let handle = Uuid::from_slice(handle.prepared_statement_handle.as_slice()) + .map_err(|e| Status::internal(format!("Error decoding handle: {}", e)))?; + let plan = self.get_plan(&handle)?; + let resp = self.execute_plan(ctx, &plan).await?; + + debug!("Responding to query..."); + Ok(resp) } + async fn get_flight_info_catalogs( &self, _query: CommandGetCatalogs, @@ -378,7 +434,7 @@ impl FlightSqlService for FlightSqlServiceImpl { "Implement get_flight_info_cross_reference", )) } - // do_get + async fn do_get_statement( &self, _ticket: TicketStatementQuery, @@ -476,21 +532,34 @@ impl FlightSqlService for FlightSqlServiceImpl { "Implement do_put_prepared_statement_update", )) } - // do_action + async fn do_action_create_prepared_statement( &self, - _query: ActionCreatePreparedStatementRequest, + query: ActionCreatePreparedStatementRequest, ) -> Result { - // TODO: implement for Flight SQL JDBC driver to work - Err(Status::unimplemented( - "Implement do_action_create_prepared_statement", - )) + let ctx = self.create_ctx().await?; + let plan = Self::prepare_statement(&query.query, &ctx).await?; + let schema_bytes = self.df_schema_to_arrow(plan.schema())?; + let handle = self.cache_plan(plan)?; + let res = ActionCreatePreparedStatementResult { + prepared_statement_handle: handle.as_bytes().to_vec(), + dataset_schema: schema_bytes, + parameter_schema: vec![] // TODO: parameters + }; + Ok(res) } + async fn do_action_close_prepared_statement( &self, - _query: ActionClosePreparedStatementRequest, + handle: ActionClosePreparedStatementRequest, ) { - unimplemented!("Implement do_action_close_prepared_statement") + let handle = Uuid::from_slice(handle.prepared_statement_handle.as_slice()); + let handle = if let Ok(handle) = handle { + handle + } else { + return; + }; + let _ = self.remove_plan(handle); } async fn register_sql_info(&self, _id: i32, _result: &SqlInfo) {} From 9e2df776974f7cb958cacb673b314552049fdb28 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 26 Jul 2022 11:21:06 -0600 Subject: [PATCH 11/11] Bot feedback --- ballista/rust/scheduler/src/flight_sql.rs | 40 +++++++++++++---------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index c7d8daa523..37704f588e 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -86,11 +86,11 @@ impl FlightSqlServiceImpl { } async fn prepare_statement( - query: &String, + query: &str, ctx: &Arc, ) -> Result { let plan = ctx - .sql(query.as_str()) + .sql(query) .await .and_then(|df| df.to_logical_plan()) .map_err(|e| Status::internal(format!("Error building plan: {}", e)))?; @@ -134,10 +134,8 @@ impl FlightSqlServiceImpl { "Error executing plan: {}", e.error )))? - }, - job_status::Status::Completed(comp) => { - Ok(Some(comp)) - }, + } + job_status::Status::Completed(comp) => Ok(Some(comp)), } } @@ -194,25 +192,34 @@ impl FlightSqlServiceImpl { fn cache_plan(&self, plan: LogicalPlan) -> Result { let handle = Uuid::new_v4(); - let mut statements = self.statements.try_lock() + let mut statements = self + .statements + .try_lock() .map_err(|e| Status::internal(format!("Error locking statements: {}", e)))?; statements.insert(handle, plan); Ok(handle) } fn get_plan(&self, handle: &Uuid) -> Result { - let statements = self.statements.try_lock() + let statements = self + .statements + .try_lock() .map_err(|e| Status::internal(format!("Error locking statements: {}", e)))?; - let plan = if let Some(plan) = statements.get(&handle) { + let plan = if let Some(plan) = statements.get(handle) { plan } else { - Err(Status::internal(format!("Statement handle not found: {}", handle)))? + Err(Status::internal(format!( + "Statement handle not found: {}", + handle + )))? }; Ok(plan.clone()) } fn remove_plan(&self, handle: Uuid) -> Result<(), Status> { - let mut statements = self.statements.try_lock() + let mut statements = self + .statements + .try_lock() .map_err(|e| Status::internal(format!("Error locking statements: {}", e)))?; statements.remove(&handle); Ok(()) @@ -292,16 +299,15 @@ impl FlightSqlServiceImpl { total_records: num_rows, total_bytes: num_bytes, }; - let resp = Response::new(info); - resp + Response::new(info) } async fn execute_plan( &self, ctx: Arc, - plan: &LogicalPlan + plan: &LogicalPlan, ) -> Result, Status> { - let job_id = self.enqueue_job(ctx, &plan).await?; + let job_id = self.enqueue_job(ctx, plan).await?; // poll for job completion let mut num_rows = 0; @@ -352,7 +358,7 @@ impl FlightSqlService for FlightSqlServiceImpl { ) -> Result, Status> { let ctx = self.create_ctx().await?; let handle = Uuid::from_slice(handle.prepared_statement_handle.as_slice()) - .map_err(|e| Status::internal(format!("Error decoding handle: {}", e)))?; + .map_err(|e| Status::internal(format!("Error decoding handle: {}", e)))?; let plan = self.get_plan(&handle)?; let resp = self.execute_plan(ctx, &plan).await?; @@ -544,7 +550,7 @@ impl FlightSqlService for FlightSqlServiceImpl { let res = ActionCreatePreparedStatementResult { prepared_statement_handle: handle.as_bytes().to_vec(), dataset_schema: schema_bytes, - parameter_schema: vec![] // TODO: parameters + parameter_schema: vec![], // TODO: parameters }; Ok(res) }