Making Spark Connect Rust more thread-safe friendly

Cargo.toml

@@ -27,7 +27,7 @@ arrow-ipc = "51.0.0"
 
 serde_json = "1.0.115"
 
-parking_lot = "0.12.1"
+parking_lot = { version="0.12.1" , features = ["send_guard"]}
 
 prost = "0.12.0"
 prost-types = "0.12.0"

examples/databricks.rs

@@ -11,9 +11,11 @@ use spark_connect_rs::{SparkSession, SparkSessionBuilder};
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let spark: SparkSession = SparkSessionBuilder::remote("sc://<workspace id>:443/;token=<personal access token>;x-databricks-cluster-id=<cluster-id>")
+    let spark:Arc<SparkSession> = Arc::new(
+        SparkSessionBuilder::remote("sc://<workspace id>:443/;token=<personal access token>;x-databricks-cluster-id=<cluster-id>")
         .build()
-        .await?;
+        .await?
+    );
 
     spark
         .range(None, 10, 1, Some(1))

examples/delta.rs

@@ -6,13 +6,15 @@
 // The remote spark session must have the spark package `io.delta:delta-spark_2.12:{DELTA_VERSION}` enabled.
 // Where the `DELTA_VERSION` is the specified Delta Lake version.
 
+use std::sync::Arc;
+
 use spark_connect_rs::{SparkSession, SparkSessionBuilder};
 
 use spark_connect_rs::dataframe::SaveMode;
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let spark: SparkSession = SparkSessionBuilder::default().build().await?;
+    let spark: Arc<SparkSession> = Arc::new(SparkSessionBuilder::default().build().await?);
 
     let paths = ["/opt/spark/examples/src/main/resources/people.csv"];
 

examples/reader.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use spark_connect_rs::{SparkSession, SparkSessionBuilder};
 
 use spark_connect_rs::functions as F;
@@ -7,7 +9,7 @@ use spark_connect_rs::functions as F;
 // printing the results as "show(...)"
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let spark: SparkSession = SparkSessionBuilder::default().build().await?;
+    let spark: Arc<SparkSession> = Arc::new(SparkSessionBuilder::default().build().await?);
 
     let path = ["/opt/spark/examples/src/main/resources/people.csv"];
 

examples/readstream.rs

@@ -3,15 +3,17 @@ use spark_connect_rs;
 use spark_connect_rs::streaming::{OutputMode, Trigger};
 use spark_connect_rs::{SparkSession, SparkSessionBuilder};
 
+use std::sync::Arc;
 use std::{thread, time};
 
 // This example demonstrates creating a Spark Stream and monitoring the progress
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let spark: SparkSession =
+    let spark: Arc<SparkSession> = Arc::new(
         SparkSessionBuilder::remote("sc://127.0.0.1:15002/;user_id=example_rs")
             .build()
-            .await?;
+            .await?,
+    );
 
     let df = spark
         .readStream()

examples/sql.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use spark_connect_rs;
 
 use spark_connect_rs::{SparkSession, SparkSessionBuilder};
@@ -7,10 +9,11 @@ use spark_connect_rs::{SparkSession, SparkSessionBuilder};
 // Displaying the results as "show(...)"
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let spark: SparkSession =
+    let spark: Arc<SparkSession> = Arc::new(
         SparkSessionBuilder::remote("sc://127.0.0.1:15002/;user_id=example_rs")
             .build()
-            .await?;
+            .await?,
+    );
 
     let df = spark
         .sql("SELECT * FROM json.`/opt/spark/examples/src/main/resources/employees.json`")

examples/writer.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use spark_connect_rs;
 
 use spark_connect_rs::{SparkSession, SparkSessionBuilder};
@@ -11,7 +13,7 @@ use spark_connect_rs::dataframe::SaveMode;
 // then reading the csv file back
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let spark: SparkSession = SparkSessionBuilder::default().build().await?;
+    let spark: Arc<SparkSession> = Arc::new(SparkSessionBuilder::default().build().await?);
 
     let df = spark
         .clone()

src/catalog.rs

@@ -1,5 +1,7 @@
 //! Spark Catalog representation through which the user may create, drop, alter or query underlying databases, tables, functions, etc.
 
+use std::sync::Arc;
+
 use arrow::array::RecordBatch;
 
 use crate::errors::SparkError;
@@ -9,11 +11,11 @@ use crate::spark;
 
 #[derive(Debug, Clone)]
 pub struct Catalog {
-    spark_session: SparkSession,
+    spark_session: Arc<SparkSession>,
 }
 
 impl Catalog {
-    pub fn new(spark_session: SparkSession) -> Self {
+    pub fn new(spark_session: Arc<SparkSession>) -> Self {
         Self { spark_session }
     }
 

src/client/mod.rs

@@ -127,15 +127,12 @@ impl ChannelBuilder {
                 channel_builder.use_ssl = true
             }
         };
-
         channel_builder.headers = Some(metadata_builder(&headers));
-
         Ok(channel_builder)
     }
 
     async fn create_client(&self) -> Result<SparkSession, Error> {
         let endpoint = format!("https://{}:{}", self.host, self.port);
-
         let channel = Endpoint::from_shared(endpoint)?.connect().await?;
 
         let service_client = SparkConnectServiceClient::with_interceptor(
@@ -413,7 +410,6 @@ where
 
         self.handle_analyze(resp)
     }
-
     fn handle_response(&mut self, resp: spark::ExecutePlanResponse) -> Result<(), SparkError> {
         self.validate_session(&resp.session_id)?;
 

src/dataframe.rs

@@ -1,5 +1,7 @@
 //! DataFrame representation for Spark Connection
 
+use std::sync::Arc;
+
 use crate::column::Column;
 use crate::errors::SparkError;
 use crate::expressions::{ToExpr, ToFilterExpr, ToVecExpr};
@@ -66,7 +68,7 @@ use arrow::util::pretty;
 #[derive(Clone, Debug)]
 pub struct DataFrame {
     /// Global [SparkSession] connecting to the remote cluster
-    pub spark_session: SparkSession,
+    pub spark_session: Arc<SparkSession>,
 
     /// Logical Plan representing the unresolved Relation
     /// which will be submitted to the remote cluster
@@ -75,7 +77,7 @@ pub struct DataFrame {
 
 impl DataFrame {
     /// create default DataFrame based on a spark session and initial logical plan
-    pub fn new(spark_session: SparkSession, logical_plan: LogicalPlanBuilder) -> DataFrame {
+    pub fn new(spark_session: Arc<SparkSession>, logical_plan: LogicalPlanBuilder) -> DataFrame {
         DataFrame {
             spark_session,
             logical_plan,
@@ -658,15 +660,11 @@ impl DataFrame {
             spark::analyze_plan_request::Analyze::Schema(spark::analyze_plan_request::Schema {
                 plan: Some(plan),
             });
-
-        let data_type = self
-            .spark_session
-            .client()
-            .analyze(schema)
-            .await?
-            .schema()?;
-
-        Ok(data_type)
+        let session = self.spark_session.clone();
+        let mut client = session.client();
+        let data_type = client.analyze(schema).await?;
+        let schema = data_type.schema()?;
+        Ok(schema.clone())
     }
 
     /// Projects a set of expressions and returns a new [DataFrame]
@@ -757,7 +755,7 @@ impl DataFrame {
     }
 
     #[allow(non_snake_case)]
-    pub fn sparkSession(self) -> SparkSession {
+    pub fn sparkSession(self) -> Arc<SparkSession> {
         self.spark_session
     }
 

src/errors.rs

@@ -79,6 +79,9 @@ impl Display for SparkError {
     }
 }
 
+unsafe impl Send for SparkError {}
+unsafe impl Sync for SparkError {}
+
 impl Error for SparkError {
     fn source(&self) -> Option<&(dyn Error + 'static)> {
         if let Self::ExternalError(e) = self {

src/readwriter.rs

@@ -1,6 +1,7 @@
 //! DataFrameReader & DataFrameWriter representations
 
 use std::collections::HashMap;
+use std::sync::Arc;
 
 use crate::errors::SparkError;
 use crate::plan::LogicalPlanBuilder;
@@ -14,14 +15,14 @@ use spark::write_operation::SaveMode;
 /// from a specific file format.
 #[derive(Clone, Debug)]
 pub struct DataFrameReader {
-    spark_session: SparkSession,
+    spark_session: Arc<SparkSession>,
     format: Option<String>,
     read_options: HashMap<String, String>,
 }
 
 impl DataFrameReader {
     /// Create a new DataFrameReader with a [SparkSession]
-    pub fn new(spark_session: SparkSession) -> Self {
+    pub fn new(spark_session: Arc<SparkSession>) -> Self {
         Self {
             spark_session,
             format: None,

src/session.rs

@@ -1,6 +1,7 @@
 //! Spark Session containing the remote gRPC client
 
 use std::collections::HashMap;
+use std::sync::Arc;
 
 use crate::catalog::Catalog;
 pub use crate::client::SparkSessionBuilder;
@@ -39,7 +40,7 @@ impl SparkSession {
     /// `end` (exclusive) with a step value `step`, and control the number
     /// of partitions with `num_partitions`
     pub fn range(
-        self,
+        self: Arc<Self>,
         start: Option<i64>,
         end: i64,
         step: i64,
@@ -55,29 +56,57 @@ impl SparkSession {
         DataFrame::new(self, LogicalPlanBuilder::from(range_relation))
     }
 
+    pub fn setCatalog(self: Arc<Self>, catalog: &str) -> DataFrame {
+        let catalog_relation = spark::relation::RelType::Catalog(spark::Catalog {
+            cat_type: Some(spark::catalog::CatType::SetCurrentCatalog(
+                spark::SetCurrentCatalog {
+                    catalog_name: catalog.to_string(),
+                },
+            )),
+        });
+
+        let logical_plan = LogicalPlanBuilder::from(catalog_relation);
+
+        DataFrame::new(self, logical_plan)
+    }
+
+    pub fn setDatabase(self: Arc<Self>, database: &str) -> DataFrame {
+        let catalog_relation = spark::relation::RelType::Catalog(spark::Catalog {
+            cat_type: Some(spark::catalog::CatType::SetCurrentDatabase(
+                spark::SetCurrentDatabase {
+                    db_name: database.to_string(),
+                },
+            )),
+        });
+
+        let logical_plan = LogicalPlanBuilder::from(catalog_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 {
+    pub fn read(self: Arc<Self>) -> DataFrameReader {
         DataFrameReader::new(self)
     }
 
     /// Returns a [DataFrameReader] that can be used to read datra in as a [DataFrame]
     #[allow(non_snake_case)]
-    pub fn readStream(self) -> DataStreamReader {
+    pub fn readStream(self: Arc<Self>) -> DataStreamReader {
         DataStreamReader::new(self)
     }
 
-    pub fn table(self, name: &str) -> Result<DataFrame, SparkError> {
+    pub fn table(self: Arc<Self>, name: &str) -> Result<DataFrame, SparkError> {
         DataFrameReader::new(self).table(name, None)
     }
 
     /// Interface through which the user may create, drop, alter or query underlying databases,
     /// tables, functions, etc.
-    pub fn catalog(self) -> Catalog {
+    pub fn catalog(self: Arc<Self>) -> Catalog {
         Catalog::new(self)
     }
 
     /// Returns a [DataFrame] representing the result of the given query
-    pub async fn sql(self, sql_query: &str) -> Result<DataFrame, SparkError> {
+    pub async fn sql(self: Arc<Self>, sql_query: &str) -> Result<DataFrame, SparkError> {
         let sql_cmd = spark::command::CommandType::SqlCommand(spark::SqlCommand {
             sql: sql_query.to_string(),
             args: HashMap::default(),
@@ -100,7 +129,7 @@ impl SparkSession {
     }
 
     #[allow(non_snake_case)]
-    pub fn createDataFrame(self, data: &RecordBatch) -> Result<DataFrame, SparkError> {
+    pub fn createDataFrame(self: Arc<Self>, data: &RecordBatch) -> Result<DataFrame, SparkError> {
         let logical_plan = LogicalPlanBuilder::local_relation(data)?;
         Ok(DataFrame::new(self, logical_plan))
     }
@@ -111,7 +140,9 @@ impl SparkSession {
     }
 
     /// Spark Connection gRPC client interface
-    pub fn client(self) -> SparkConnectClient<InterceptedService<Channel, MetadataInterceptor>> {
-        self.client
+    pub fn client(
+        self: Arc<Self>,
+    ) -> SparkConnectClient<InterceptedService<Channel, MetadataInterceptor>> {
+        self.client.clone()
     }
 }