src/execution/context.rs - spark-connect-rust - Git at Google

 use std::collections::HashMap;
 use std::io::Read;
 use std::sync::Arc;

 use crate::dataframe::{DataFrame, DataFrameReader};
 use crate::plan::LogicalPlanBuilder;
 use crate::spark;

 use arrow::error::ArrowError;
 use arrow::record_batch::RecordBatch;
 use arrow_ipc::reader::StreamReader;

 use spark::execute_plan_response::{ArrowBatch, Metrics};
 use spark::expression::Literal;
 use spark::spark_connect_service_client::SparkConnectServiceClient;
 use spark::{DataType, ExecutePlanResponse};

 use tokio::sync::Mutex;
 use tonic::transport::{Channel, Error};
 use tonic::Streaming;

 use url::Url;
 use uuid::Uuid;

 /// ChannelBuilder validates a connection string
 /// based on the requirements from [Spark Documentation](https://github.com/apache/spark/blob/master/connector/connect/docs/client-connection-string.md)
 #[derive(Clone, Debug)]
 struct ChannelBuilder {
     host: String,
     port: u16,
     token: Option<String>,
     user: Option<String>,
     headers: Option<HashMap<String, String>>,
 }

 impl Default for ChannelBuilder {
     fn default() -> Self {
         ChannelBuilder::build("sc://127.0.0.1:15002".to_string()).unwrap()
     }
 }

 impl ChannelBuilder {
     /// Build and Validate a connnection string
     pub fn build(connection: String) -> Result<ChannelBuilder, String> {
         let url =
             Url::parse(connection.as_str()).map_err(|_| "Failed to parse the url.".to_string())?;

         if url.scheme() != "sc" {
             return Err("Scheme is not set to 'sc'".to_string());
         };

         let host = url
             .host_str()
             .ok_or("Missing host in the URL.".to_string())?
             .to_string();

         let port = url.port().ok_or("Missing port in the URL.".to_string())?;

         let mut channel_builder = ChannelBuilder {
             host,
             port,
             token: None,
             user: None,
             headers: None,
         };

         let path: Vec<&str> = url.path().split(';').collect();

         if path.is_empty() || (path.len() == 1 && (path[0].is_empty() || path[0] == "/")) {
             return Ok(channel_builder);
         }

         let mut headers: HashMap<String, String> = path
             .into_iter()
             .filter(|&pair| (pair != "/") & (!pair.is_empty()))
             .map(|pair| {
                 let mut parts = pair.splitn(2, '=');
                 (
                     parts.next().unwrap_or("").to_string(),
                     parts.next().unwrap_or("").to_string(),
                 )
             })
             .collect();

         if headers.is_empty() {
             return Ok(channel_builder);
         }

         channel_builder.token = headers.remove("token");
         channel_builder.user = headers.remove("user_id");
         channel_builder.headers = Some(headers);

         Ok(channel_builder)
     }
 }

 /// SparkSessionBuilder creates a remote Spark Session a connection string.
 ///
 /// The connection string is define based on the requirements from [Spark Documentation](https://github.com/apache/spark/blob/master/connector/connect/docs/client-connection-string.md)
 #[derive(Clone, Debug)]
 pub struct SparkSessionBuilder {
     channel_builder: ChannelBuilder,
 }

 /// Default connects a Spark cluster running at `sc://127.0.0.1:15002/`
 impl Default for SparkSessionBuilder {
     fn default() -> Self {
         let channel_builder = ChannelBuilder::default();

         Self { channel_builder }
     }
 }

 impl SparkSessionBuilder {
     fn new(connection: String) -> Self {
         let channel_builder = ChannelBuilder::build(connection).unwrap();

         Self { channel_builder }
     }

     /// Validate a connect string for a remote Spark Session
     ///
     /// String must conform to the [Spark Documentation](https://github.com/apache/spark/blob/master/connector/connect/docs/client-connection-string.md)
     pub fn remote(connection: String) -> Self {
         Self::new(connection)
     }

     /// Attempt to connect to a remote Spark Session
     ///
     /// and return a [SparkSession]
     pub async fn build(self) -> Result<SparkSession, Error> {
         let url = format!(
             "https://{}:{}",
             self.channel_builder.host, self.channel_builder.port
         );

         let client = Arc::new(Mutex::new(
             SparkConnectServiceClient::connect(url.clone()).await?,
         ));

         Ok(SparkSession {
             client,
             session_id: Uuid::new_v4().to_string(),
             metadata: self.channel_builder.headers,
             user: self.channel_builder.user,
             token: self.channel_builder.token,
         })
     }
 }

 /// The entry point to connecting to a Spark Cluster
 /// using the Spark Connection gRPC protocol.
 #[allow(dead_code)]
 #[derive(Clone, Debug)]
 pub struct SparkSession {
     /// Spark Connection gRPC client interface
     pub client: Arc<Mutex<SparkConnectServiceClient<Channel>>>,

     /// Spark Session ID
     pub session_id: String,

     /// gRPC metadata collected from the connection string
     pub metadata: Option<HashMap<String, String>>,
     user: Option<String>,

     token: Option<String>,
 }

 impl SparkSession {
     /// Create a [DataFrame] with a spingle column named `id`,
     /// containing elements in a range from `start` (default 0) to
     /// `end` (exclusive) with a step value `step`, and control the number
     /// of partitions with `num_partitions`
     pub fn range(
         self,
         start: Option<i64>,
         end: i64,
         step: i64,
         num_partitions: Option<i32>,
     ) -> DataFrame {
         let range_relation = spark::Relation {
             common: Some(spark::RelationCommon {
                 source_info: "na".to_string(),
                 plan_id: Some(1),
             }),
             rel_type: Some(spark::relation::RelType::Range(spark::Range {
                 start,
                 end,
                 step,
                 num_partitions,
             })),
         };

         let logical_plan = LogicalPlanBuilder::new(range_relation);

         DataFrame::new(self, logical_plan)
     }

     /// Returns a [DataFrameReader] that can be used to read datra in as a [DataFrame]
     pub fn read(self) -> DataFrameReader {
         DataFrameReader::new(self)
     }

     /// Returns a [DataFrame] representing the result of the given query
     pub fn sql(self, sql_query: &str) -> DataFrame {
         let kind = Some(spark::data_type::Kind::Null(spark::data_type::Null {
             type_variation_reference: 1,
         }));

         let sql_command = spark::Relation {
             common: Some(spark::RelationCommon {
                 source_info: "NA".to_string(),
                 plan_id: Some(1),
             }),
             rel_type: Some(spark::relation::RelType::Sql(spark::Sql {
                 query: sql_query.to_string(),
                 args: HashMap::new(),
                 pos_args: vec![Literal {
                     literal_type: Some(spark::expression::literal::LiteralType::Null(DataType {
                         kind,
                     })),
                 }],
             })),
         };

         let logical_plan = LogicalPlanBuilder::new(sql_command);

         DataFrame::new(self, logical_plan)
     }

     fn build_execute_plan_request(&self, plan: Option<spark::Plan>) -> spark::ExecutePlanRequest {
         spark::ExecutePlanRequest {
             session_id: self.session_id.clone(),
             user_context: Some(spark::UserContext {
                 user_id: self.user.clone().unwrap_or("NA".to_string()),
                 user_name: self.user.clone().unwrap_or("NA".to_string()),
                 extensions: vec![],
             }),
             operation_id: None,
             plan,
             client_type: Some("_SPARK_CONNECT_RUST".to_string()),
             request_options: vec![],
             tags: vec![],
         }
     }

     fn build_analyze_plan_request(
         &self,
         analyze: Option<spark::analyze_plan_request::Analyze>,
     ) -> spark::AnalyzePlanRequest {
         spark::AnalyzePlanRequest {
             session_id: self.session_id.clone(),
             user_context: Some(spark::UserContext {
                 user_id: self.user.clone().unwrap_or("NA".to_string()),
                 user_name: self.user.clone().unwrap_or("NA".to_string()),
                 extensions: vec![],
             }),
             client_type: Some("_SPARK_CONNECT_RUST".to_string()),
             analyze,
         }
     }

     async fn execute_plan(
         &mut self,
         plan: Option<spark::Plan>,
     ) -> Result<Streaming<ExecutePlanResponse>, tonic::Status> {
         let exc_plan = self.build_execute_plan_request(plan);

         let mut client = self.client.lock().await;

         let value = client.execute_plan(exc_plan).await?.into_inner();

         Ok(value)
     }

     /// Call a service on the remote Spark Connect server by running
     /// a provided [spark::Plan].
     ///
     /// A [spark::Plan] produces a vector of [RecordBatch] records
     pub async fn consume_plan(
         &mut self,
         plan: Option<spark::Plan>,
     ) -> Result<Vec<RecordBatch>, ArrowError> {
         let mut stream = self
             .execute_plan(plan)
             .await
             .map_err(|err| ArrowError::IoError(err.to_string()))?;

         let mut handler = ResponseHandler::new();

         while let Some(resp) = stream
             .message()
             .await
             .map_err(|err| ArrowError::IoError(err.to_string()))?
         {
             let _ = handler.handle_response(&resp);
         }
         handler.records()
     }

     pub async fn analyze_plan(
         &mut self,
         analyze: Option<spark::analyze_plan_request::Analyze>,
     ) -> spark::analyze_plan_response::Result {
         let request = self.build_analyze_plan_request(analyze);
         let mut client = self.client.lock().await;

         let stream = client.analyze_plan(request).await.unwrap().into_inner();

         stream.result.unwrap()
     }
 }

 struct ResponseHandler {
     schema: Option<DataType>,
     data: Vec<Option<ArrowBatch>>,
     metrics: Option<Metrics>,
 }

 impl Default for ResponseHandler {
     fn default() -> Self {
         Self::new()
     }
 }

 impl ResponseHandler {
     fn new() -> ResponseHandler {
         ResponseHandler {
             schema: None,
             data: vec![],
             metrics: None,
         }
     }

     fn handle_response(&mut self, response: &ExecutePlanResponse) -> Result<(), String> {
         if let Some(schema) = response.schema.as_ref() {
             self.schema = Some(schema.clone());
         }
         if let Some(metrics) = response.metrics.as_ref() {
             self.metrics = Some(metrics.clone());
         }
         if let Some(data) = response.response_type.as_ref() {
             match data {
                 spark::execute_plan_response::ResponseType::ArrowBatch(batch) => {
                     self.data.push(Some(batch.clone()));
                 }
                 _ => {
                     return Err("Not implemented".to_string());
                 }
             }
         }
         Ok(())
     }

     fn records(self) -> Result<Vec<RecordBatch>, ArrowError> {
         let mut accumulator: Vec<Vec<RecordBatch>> = vec![vec![]];
         for batch in self.data.into_iter().flatten() {
             accumulator.push(deserialize(batch)?);
         }

         Ok(accumulator
             .into_iter()
             .flatten()
             .collect::<Vec<RecordBatch>>())
     }
 }

 struct ArrowBatchReader {
     batch: ArrowBatch,
 }

 impl Read for ArrowBatchReader {
     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
         Read::read(&mut self.batch.data.as_slice(), buf)
     }
 }

 fn deserialize(batch: ArrowBatch) -> Result<Vec<RecordBatch>, ArrowError> {
     let wrapper = ArrowBatchReader { batch };
     let reader = StreamReader::try_new(wrapper, None)?;
     let mut rows = Vec::new();
     for record in reader {
         rows.push(record?)
     }
     Ok(rows)
 }

 #[cfg(test)]
 mod tests {
     use super::*;

     #[test]
     fn test_channel_builder_default() {
         let expected_url = "127.0.0.1:15002".to_string();

         let cb = ChannelBuilder::default();

         let output_url = format!("{}:{}", cb.host, cb.port);

         assert_eq!(expected_url, output_url)
     }

     #[test]
     #[should_panic(expected = "Scheme is not set to 'sc")]
     fn test_panic_incorrect_url_scheme() {
         let connection = "http://127.0.0.1:15002".to_string();

         ChannelBuilder::build(connection).unwrap();
     }

     #[test]
     #[should_panic(expected = "Failed to parse the url.")]
     fn test_panic_missing_url_host() {
         let connection = "sc://:15002".to_string();

         ChannelBuilder::build(connection).unwrap();
     }

     #[test]
     #[should_panic(expected = "Missing port in the URL")]
     fn test_panic_missing_url_port() {
         let connection = "sc://127.0.0.1".to_string();

         ChannelBuilder::build(connection).unwrap();
     }

     #[test]
     fn test_spark_session_builder() {
         let connection =
             "sc://myhost.com:443/;use_ssl=true;token=ABCDEFG;user_agent=some_agent;user_id=user123"
                 .to_string();

         let ssbuilder = SparkSessionBuilder::remote(connection);

         assert_eq!("myhost.com".to_string(), ssbuilder.channel_builder.host);
         assert_eq!(443, ssbuilder.channel_builder.port);
         assert_eq!(
             "ABCDEFG".to_string(),
             ssbuilder.channel_builder.token.unwrap()
         );
         assert_eq!(
             "user123".to_string(),
             ssbuilder.channel_builder.user.unwrap()
         );
         assert_eq!(
             Some(&"true".to_string()),
             ssbuilder
                 .channel_builder
                 .headers
                 .clone()
                 .unwrap()
                 .get("use_ssl")
         );
         assert_eq!(
             Some(&"some_agent".to_string()),
             ssbuilder
                 .channel_builder
                 .headers
                 .clone()
                 .unwrap()
                 .get("user_agent")
         );
     }

     #[tokio::test]
     async fn test_spark_session_create() {
         let connection = "sc://localhost:15002/;use_ssl=true;token=ABCDEFG;user_agent=some_agent;user_id=user123".to_string();

         let spark = SparkSessionBuilder::remote(connection).build().await;

         assert!(spark.is_ok());
         assert_eq!(
             Some(&"true".to_string()),
             spark.unwrap().metadata.unwrap().get("use_ssl")
         );
     }
 }
	use std::collections::HashMap;
	use std::io::Read;
	use std::sync::Arc;

	use crate::dataframe::{DataFrame, DataFrameReader};
	use crate::plan::LogicalPlanBuilder;
	use crate::spark;

	use arrow::error::ArrowError;
	use arrow::record_batch::RecordBatch;
	use arrow_ipc::reader::StreamReader;

	use spark::execute_plan_response::{ArrowBatch, Metrics};
	use spark::expression::Literal;
	use spark::spark_connect_service_client::SparkConnectServiceClient;
	use spark::{DataType, ExecutePlanResponse};

	use tokio::sync::Mutex;
	use tonic::transport::{Channel, Error};
	use tonic::Streaming;

	use url::Url;
	use uuid::Uuid;

	/// ChannelBuilder validates a connection string
	/// based on the requirements from [Spark Documentation](https://github.com/apache/spark/blob/master/connector/connect/docs/client-connection-string.md)
	#[derive(Clone, Debug)]
	struct ChannelBuilder {
	host: String,
	port: u16,
	token: Option<String>,
	user: Option<String>,
	headers: Option<HashMap<String, String>>,
	}

	impl Default for ChannelBuilder {
	fn default() -> Self {
	ChannelBuilder::build("sc://127.0.0.1:15002".to_string()).unwrap()
	}
	}

	impl ChannelBuilder {
	/// Build and Validate a connnection string
	pub fn build(connection: String) -> Result<ChannelBuilder, String> {
	let url =
	Url::parse(connection.as_str()).map_err(\|_\| "Failed to parse the url.".to_string())?;

	if url.scheme() != "sc" {
	return Err("Scheme is not set to 'sc'".to_string());
	};

	let host = url
	.host_str()
	.ok_or("Missing host in the URL.".to_string())?
	.to_string();

	let port = url.port().ok_or("Missing port in the URL.".to_string())?;

	let mut channel_builder = ChannelBuilder {
	host,
	port,
	token: None,
	user: None,
	headers: None,
	};

	let path: Vec<&str> = url.path().split(';').collect();

	if path.is_empty() \|\| (path.len() == 1 && (path[0].is_empty() \|\| path[0] == "/")) {
	return Ok(channel_builder);
	}

	let mut headers: HashMap<String, String> = path
	.into_iter()
	.filter(\|&pair\| (pair != "/") & (!pair.is_empty()))
	.map(\|pair\| {
	let mut parts = pair.splitn(2, '=');
	(
	parts.next().unwrap_or("").to_string(),
	parts.next().unwrap_or("").to_string(),
	)
	})
	.collect();

	if headers.is_empty() {
	return Ok(channel_builder);
	}

	channel_builder.token = headers.remove("token");
	channel_builder.user = headers.remove("user_id");
	channel_builder.headers = Some(headers);

	Ok(channel_builder)
	}
	}

	/// SparkSessionBuilder creates a remote Spark Session a connection string.
	///
	/// The connection string is define based on the requirements from [Spark Documentation](https://github.com/apache/spark/blob/master/connector/connect/docs/client-connection-string.md)
	#[derive(Clone, Debug)]
	pub struct SparkSessionBuilder {
	channel_builder: ChannelBuilder,
	}

	/// Default connects a Spark cluster running at `sc://127.0.0.1:15002/`
	impl Default for SparkSessionBuilder {
	fn default() -> Self {
	let channel_builder = ChannelBuilder::default();

	Self { channel_builder }
	}
	}

	impl SparkSessionBuilder {
	fn new(connection: String) -> Self {
	let channel_builder = ChannelBuilder::build(connection).unwrap();

	Self { channel_builder }
	}

	/// Validate a connect string for a remote Spark Session
	///
	/// String must conform to the [Spark Documentation](https://github.com/apache/spark/blob/master/connector/connect/docs/client-connection-string.md)
	pub fn remote(connection: String) -> Self {
	Self::new(connection)
	}

	/// Attempt to connect to a remote Spark Session
	///
	/// and return a [SparkSession]
	pub async fn build(self) -> Result<SparkSession, Error> {
	let url = format!(
	"https://{}:{}",
	self.channel_builder.host, self.channel_builder.port
	);

	let client = Arc::new(Mutex::new(
	SparkConnectServiceClient::connect(url.clone()).await?,
	));

	Ok(SparkSession {
	client,
	session_id: Uuid::new_v4().to_string(),
	metadata: self.channel_builder.headers,
	user: self.channel_builder.user,
	token: self.channel_builder.token,
	})
	}
	}

	/// The entry point to connecting to a Spark Cluster
	/// using the Spark Connection gRPC protocol.
	#[allow(dead_code)]
	#[derive(Clone, Debug)]
	pub struct SparkSession {
	/// Spark Connection gRPC client interface
	pub client: Arc<Mutex<SparkConnectServiceClient<Channel>>>,

	/// Spark Session ID
	pub session_id: String,

	/// gRPC metadata collected from the connection string
	pub metadata: Option<HashMap<String, String>>,
	user: Option<String>,

	token: Option<String>,
	}

	impl SparkSession {
	/// Create a [DataFrame] with a spingle column named `id`,
	/// containing elements in a range from `start` (default 0) to
	/// `end` (exclusive) with a step value `step`, and control the number
	/// of partitions with `num_partitions`
	pub fn range(
	self,
	start: Option<i64>,
	end: i64,
	step: i64,
	num_partitions: Option<i32>,
	) -> DataFrame {
	let range_relation = spark::Relation {
	common: Some(spark::RelationCommon {
	source_info: "na".to_string(),
	plan_id: Some(1),
	}),
	rel_type: Some(spark::relation::RelType::Range(spark::Range {
	start,
	end,
	step,
	num_partitions,
	})),
	};

	let logical_plan = LogicalPlanBuilder::new(range_relation);

	DataFrame::new(self, logical_plan)
	}

	/// Returns a [DataFrameReader] that can be used to read datra in as a [DataFrame]
	pub fn read(self) -> DataFrameReader {
	DataFrameReader::new(self)
	}

	/// Returns a [DataFrame] representing the result of the given query
	pub fn sql(self, sql_query: &str) -> DataFrame {
	let kind = Some(spark::data_type::Kind::Null(spark::data_type::Null {
	type_variation_reference: 1,
	}));

	let sql_command = spark::Relation {
	common: Some(spark::RelationCommon {
	source_info: "NA".to_string(),
	plan_id: Some(1),
	}),
	rel_type: Some(spark::relation::RelType::Sql(spark::Sql {
	query: sql_query.to_string(),
	args: HashMap::new(),
	pos_args: vec![Literal {
	literal_type: Some(spark::expression::literal::LiteralType::Null(DataType {
	kind,
	})),
	}],
	})),
	};

	let logical_plan = LogicalPlanBuilder::new(sql_command);

	DataFrame::new(self, logical_plan)
	}

	fn build_execute_plan_request(&self, plan: Option<spark::Plan>) -> spark::ExecutePlanRequest {
	spark::ExecutePlanRequest {
	session_id: self.session_id.clone(),
	user_context: Some(spark::UserContext {
	user_id: self.user.clone().unwrap_or("NA".to_string()),
	user_name: self.user.clone().unwrap_or("NA".to_string()),
	extensions: vec![],
	}),
	operation_id: None,
	plan,
	client_type: Some("_SPARK_CONNECT_RUST".to_string()),
	request_options: vec![],
	tags: vec![],
	}
	}

	fn build_analyze_plan_request(
	&self,
	analyze: Option<spark::analyze_plan_request::Analyze>,
	) -> spark::AnalyzePlanRequest {
	spark::AnalyzePlanRequest {
	session_id: self.session_id.clone(),
	user_context: Some(spark::UserContext {
	user_id: self.user.clone().unwrap_or("NA".to_string()),
	user_name: self.user.clone().unwrap_or("NA".to_string()),
	extensions: vec![],
	}),
	client_type: Some("_SPARK_CONNECT_RUST".to_string()),
	analyze,
	}
	}

	async fn execute_plan(
	&mut self,
	plan: Option<spark::Plan>,
	) -> Result<Streaming<ExecutePlanResponse>, tonic::Status> {
	let exc_plan = self.build_execute_plan_request(plan);

	let mut client = self.client.lock().await;

	let value = client.execute_plan(exc_plan).await?.into_inner();

	Ok(value)
	}

	/// Call a service on the remote Spark Connect server by running
	/// a provided [spark::Plan].
	///
	/// A [spark::Plan] produces a vector of [RecordBatch] records
	pub async fn consume_plan(
	&mut self,
	plan: Option<spark::Plan>,
	) -> Result<Vec<RecordBatch>, ArrowError> {
	let mut stream = self
	.execute_plan(plan)
	.await
	.map_err(\|err\| ArrowError::IoError(err.to_string()))?;

	let mut handler = ResponseHandler::new();

	while let Some(resp) = stream
	.message()
	.await
	.map_err(\|err\| ArrowError::IoError(err.to_string()))?
	{
	let _ = handler.handle_response(&resp);
	}
	handler.records()
	}

	pub async fn analyze_plan(
	&mut self,
	analyze: Option<spark::analyze_plan_request::Analyze>,
	) -> spark::analyze_plan_response::Result {
	let request = self.build_analyze_plan_request(analyze);
	let mut client = self.client.lock().await;

	let stream = client.analyze_plan(request).await.unwrap().into_inner();

	stream.result.unwrap()
	}
	}

	struct ResponseHandler {
	schema: Option<DataType>,
	data: Vec<Option<ArrowBatch>>,
	metrics: Option<Metrics>,
	}

	impl Default for ResponseHandler {
	fn default() -> Self {
	Self::new()
	}
	}

	impl ResponseHandler {
	fn new() -> ResponseHandler {
	ResponseHandler {
	schema: None,
	data: vec![],
	metrics: None,
	}
	}

	fn handle_response(&mut self, response: &ExecutePlanResponse) -> Result<(), String> {
	if let Some(schema) = response.schema.as_ref() {
	self.schema = Some(schema.clone());
	}
	if let Some(metrics) = response.metrics.as_ref() {
	self.metrics = Some(metrics.clone());
	}
	if let Some(data) = response.response_type.as_ref() {
	match data {
	spark::execute_plan_response::ResponseType::ArrowBatch(batch) => {
	self.data.push(Some(batch.clone()));
	}
	_ => {
	return Err("Not implemented".to_string());
	}
	}
	}
	Ok(())
	}

	fn records(self) -> Result<Vec<RecordBatch>, ArrowError> {
	let mut accumulator: Vec<Vec<RecordBatch>> = vec![vec![]];
	for batch in self.data.into_iter().flatten() {
	accumulator.push(deserialize(batch)?);
	}

	Ok(accumulator
	.into_iter()
	.flatten()
	.collect::<Vec<RecordBatch>>())
	}
	}

	struct ArrowBatchReader {
	batch: ArrowBatch,
	}

	impl Read for ArrowBatchReader {
	fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
	Read::read(&mut self.batch.data.as_slice(), buf)
	}
	}

	fn deserialize(batch: ArrowBatch) -> Result<Vec<RecordBatch>, ArrowError> {
	let wrapper = ArrowBatchReader { batch };
	let reader = StreamReader::try_new(wrapper, None)?;
	let mut rows = Vec::new();
	for record in reader {
	rows.push(record?)
	}
	Ok(rows)
	}

	#[cfg(test)]
	mod tests {
	use super::*;

	#[test]
	fn test_channel_builder_default() {
	let expected_url = "127.0.0.1:15002".to_string();

	let cb = ChannelBuilder::default();

	let output_url = format!("{}:{}", cb.host, cb.port);

	assert_eq!(expected_url, output_url)
	}

	#[test]
	#[should_panic(expected = "Scheme is not set to 'sc")]
	fn test_panic_incorrect_url_scheme() {
	let connection = "http://127.0.0.1:15002".to_string();

	ChannelBuilder::build(connection).unwrap();
	}

	#[test]
	#[should_panic(expected = "Failed to parse the url.")]
	fn test_panic_missing_url_host() {
	let connection = "sc://:15002".to_string();

	ChannelBuilder::build(connection).unwrap();
	}

	#[test]
	#[should_panic(expected = "Missing port in the URL")]
	fn test_panic_missing_url_port() {
	let connection = "sc://127.0.0.1".to_string();

	ChannelBuilder::build(connection).unwrap();
	}

	#[test]
	fn test_spark_session_builder() {
	let connection =
	"sc://myhost.com:443/;use_ssl=true;token=ABCDEFG;user_agent=some_agent;user_id=user123"
	.to_string();

	let ssbuilder = SparkSessionBuilder::remote(connection);

	assert_eq!("myhost.com".to_string(), ssbuilder.channel_builder.host);
	assert_eq!(443, ssbuilder.channel_builder.port);
	assert_eq!(
	"ABCDEFG".to_string(),
	ssbuilder.channel_builder.token.unwrap()
	);
	assert_eq!(
	"user123".to_string(),
	ssbuilder.channel_builder.user.unwrap()
	);
	assert_eq!(
	Some(&"true".to_string()),
	ssbuilder
	.channel_builder
	.headers
	.clone()
	.unwrap()
	.get("use_ssl")
	);
	assert_eq!(
	Some(&"some_agent".to_string()),
	ssbuilder
	.channel_builder
	.headers
	.clone()
	.unwrap()
	.get("user_agent")
	);
	}

	#[tokio::test]
	async fn test_spark_session_create() {
	let connection = "sc://localhost:15002/;use_ssl=true;token=ABCDEFG;user_agent=some_agent;user_id=user123".to_string();

	let spark = SparkSessionBuilder::remote(connection).build().await;

	assert!(spark.is_ok());
	assert_eq!(
	Some(&"true".to_string()),
	spark.unwrap().metadata.unwrap().get("use_ssl")
	);
	}
	}