feat(generics): cleaner generics

- add better generic types to traits
- add new trait implementation for an array of expressions
- update example

examples/delta.rs

@@ -14,7 +14,7 @@ use spark_connect_rs::dataframe::SaveMode;
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let mut spark: SparkSession = SparkSessionBuilder::default().build().await?;
 
-    let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];
+    let paths = ["/opt/spark/examples/src/main/resources/people.csv"];
 
     let df = spark
         .clone()

examples/reader.rs

@@ -9,16 +9,16 @@ use spark_connect_rs::functions as F;
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let spark: SparkSession = SparkSessionBuilder::default().build().await?;
 
-    let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];
+    let path = ["/opt/spark/examples/src/main/resources/people.csv"];
 
     let mut df = spark
         .read()
         .format("csv")
         .option("header", "True")
         .option("delimiter", ";")
-        .load(paths);
+        .load(path);
 
-    df.select(vec![
+    df.select([
         F::col("name"),
         F::col("age").cast("int").alias("age_int"),
         (F::lit(3.0) + F::col("age").cast("int")).alias("addition"),

examples/writer.rs

@@ -32,7 +32,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .read()
         .format("csv")
         .option("header", "true")
-        .load(vec![path.to_string()]);
+        .load([path]);
 
     df.show(Some(10), None, None).await?;
 

src/catalog.rs

@@ -54,7 +54,9 @@ impl Catalog {
 
     /// Returns a list of catalogs in this session
     #[allow(non_snake_case)]
-    pub async fn listCatalogs(&mut self, pattern: Option<String>) -> Vec<RecordBatch> {
+    pub async fn listCatalogs(&mut self, pattern: Option<&str>) -> Vec<RecordBatch> {
+        let pattern = pattern.map(|val| val.to_owned());
+
         let cat_type = Some(spark::catalog::CatType::ListCatalogs(spark::ListCatalogs {
             pattern,
         }));
@@ -72,7 +74,9 @@ impl Catalog {
 
     /// Returns a list of databases in this session
     #[allow(non_snake_case)]
-    pub async fn listDatabases(&mut self, pattern: Option<String>) -> Vec<RecordBatch> {
+    pub async fn listDatabases(&mut self, pattern: Option<&str>) -> Vec<RecordBatch> {
+        let pattern = pattern.map(|val| val.to_owned());
+
         let cat_type = Some(spark::catalog::CatType::ListDatabases(
             spark::ListDatabases { pattern },
         ));
@@ -92,12 +96,12 @@ impl Catalog {
     #[allow(non_snake_case)]
     pub async fn listTables(
         &mut self,
-        dbName: Option<String>,
-        pattern: Option<String>,
+        dbName: Option<&str>,
+        pattern: Option<&str>,
     ) -> Vec<RecordBatch> {
         let cat_type = Some(spark::catalog::CatType::ListTables(spark::ListTables {
-            db_name: dbName,
-            pattern,
+            db_name: dbName.map(|db| db.to_owned()),
+            pattern: pattern.map(|val| val.to_owned()),
         }));
 
         let rel_type = spark::relation::RelType::Catalog(spark::Catalog { cat_type });
@@ -113,14 +117,10 @@ impl Catalog {
 
     /// Returns a list of columns for the given tables/views in the specific database
     #[allow(non_snake_case)]
-    pub async fn listColumns(
-        &mut self,
-        tableName: String,
-        dbName: Option<String>,
-    ) -> Vec<RecordBatch> {
+    pub async fn listColumns(&mut self, tableName: &str, dbName: Option<&str>) -> Vec<RecordBatch> {
         let cat_type = Some(spark::catalog::CatType::ListColumns(spark::ListColumns {
-            table_name: tableName,
-            db_name: dbName,
+            table_name: tableName.to_owned(),
+            db_name: dbName.map(|val| val.to_owned()),
         }));
 
         let rel_type = spark::relation::RelType::Catalog(spark::Catalog { cat_type });
@@ -207,10 +207,7 @@ mod tests {
             .await
             .unwrap();
 
-        let value = spark
-            .catalog()
-            .listDatabases(Some("*rust".to_string()))
-            .await;
+        let value = spark.catalog().listDatabases(Some("*rust")).await;
 
         assert_eq!(4, value[0].num_columns());
         assert_eq!(1, value[0].num_rows());

src/dataframe.rs

@@ -1,7 +1,5 @@
 //! DataFrame with Reader/Writer repesentation
 
-use std::collections::HashMap;
-
 use crate::column::Column;
 use crate::errors::SparkError;
 use crate::expressions::{ToExpr, ToFilterExpr, ToVecExpr};
@@ -159,44 +157,36 @@ impl DataFrame {
         Some(data?.value(0))
     }
 
-    #[allow(non_snake_case, dead_code)]
+    #[allow(non_snake_case)]
     pub async fn createTempView(&mut self, name: &str) {
-        self.create_view_cmd(name.to_string(), false, false)
-            .await
-            .unwrap()
+        self.create_view_cmd(name, false, false).await.unwrap()
     }
 
-    #[allow(non_snake_case, dead_code)]
+    #[allow(non_snake_case)]
     pub async fn createGlobalTempView(&mut self, name: &str) {
-        self.create_view_cmd(name.to_string(), true, false)
-            .await
-            .unwrap()
+        self.create_view_cmd(name, true, false).await.unwrap()
     }
 
-    #[allow(non_snake_case, dead_code)]
+    #[allow(non_snake_case)]
     pub async fn createOrReplaceGlobalTempView(&mut self, name: &str) {
-        self.create_view_cmd(name.to_string(), true, true)
-            .await
-            .unwrap()
+        self.create_view_cmd(name, true, true).await.unwrap()
     }
 
-    #[allow(non_snake_case, dead_code)]
+    #[allow(non_snake_case)]
     pub async fn createOrReplaceTempView(&mut self, name: &str) {
-        self.create_view_cmd(name.to_string(), false, true)
-            .await
-            .unwrap()
+        self.create_view_cmd(name, false, true).await.unwrap()
     }
 
     async fn create_view_cmd(
         &mut self,
-        name: String,
+        name: &str,
         is_global: bool,
         replace: bool,
     ) -> Result<(), SparkError> {
         let command_type =
             spark::command::CommandType::CreateDataframeView(spark::CreateDataFrameViewCommand {
                 input: Some(self.logical_plan.clone().relation),
-                name,
+                name: name.to_string(),
                 is_global,
                 replace,
             });
@@ -357,7 +347,10 @@ impl DataFrame {
     }
 
     #[allow(non_snake_case)]
-    pub fn freqItems(&mut self, cols: Vec<&str>, support: Option<f64>) -> DataFrame {
+    pub fn freqItems<'a, I>(&mut self, cols: I, support: Option<f64>) -> DataFrame
+    where
+        I: IntoIterator<Item = &'a str>,
+    {
         DataFrame::new(
             self.spark_session.clone(),
             self.logical_plan.freqItems(cols, support),
@@ -455,7 +448,10 @@ impl DataFrame {
     }
 
     #[allow(non_snake_case)]
-    pub fn orderBy(&mut self, cols: Vec<Column>) -> DataFrame {
+    pub fn orderBy<I>(&mut self, cols: I) -> DataFrame
+    where
+        I: IntoIterator<Item = Column>,
+    {
         DataFrame::new(self.spark_session.clone(), self.logical_plan.sort(cols))
     }
 
@@ -583,7 +579,10 @@ impl DataFrame {
     /// }
     /// ```
     #[allow(non_snake_case)]
-    pub fn selectExpr(&mut self, cols: Vec<&str>) -> DataFrame {
+    pub fn selectExpr<'a, I>(&mut self, cols: I) -> DataFrame
+    where
+        I: IntoIterator<Item = &'a str>,
+    {
         DataFrame::new(
             self.spark_session.clone(),
             self.logical_plan.select_expr(cols),
@@ -692,7 +691,10 @@ impl DataFrame {
     }
 
     #[allow(non_snake_case)]
-    pub fn toDF(&mut self, cols: Vec<&str>) -> DataFrame {
+    pub fn toDF<'a, I>(&mut self, cols: I) -> DataFrame
+    where
+        I: IntoIterator<Item = &'a str>,
+    {
         DataFrame::new(self.spark_session.clone(), self.logical_plan.to_df(cols))
     }
 
@@ -752,18 +754,27 @@ impl DataFrame {
     }
 
     #[allow(non_snake_case)]
-    pub fn withColumns(&mut self, colMap: HashMap<&str, Column>) -> DataFrame {
+    pub fn withColumns<I, K>(&mut self, colMap: I) -> DataFrame
+    where
+        I: IntoIterator<Item = (K, Column)>,
+        K: ToString,
+    {
         DataFrame::new(
             self.spark_session.clone(),
             self.logical_plan.withColumns(colMap),
         )
     }
 
     /// Returns a new [DataFrame] by renaming multiple columns from a
-    /// `HashMap<String, String>` containing the `existing` as the key
-    /// and the `new` as the value.
+    /// an iterator of containing a key/value pair with the key as the `existing`
+    /// column name and the value as the `new` column name.
     #[allow(non_snake_case)]
-    pub fn withColumnsRenamed(&mut self, cols: HashMap<String, String>) -> DataFrame {
+    pub fn withColumnsRenamed<I, K, V>(&mut self, cols: I) -> DataFrame
+    where
+        I: IntoIterator<Item = (K, V)>,
+        K: AsRef<str>,
+        V: AsRef<str>,
+    {
         DataFrame::new(
             self.spark_session.clone(),
             self.logical_plan.withColumnsRenamed(cols),
@@ -903,14 +914,14 @@ mod tests {
             cols
         );
 
-        let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];
+        let path = ["/opt/spark/examples/src/main/resources/people.csv"];
 
         let cols = spark
             .read()
             .format("csv")
             .option("header", "True")
             .option("delimiter", ";")
-            .load(paths)
+            .load(path)
             .columns()
             .await;
 
@@ -1000,14 +1011,14 @@ mod tests {
     async fn test_describe() {
         let spark = setup().await;
 
-        let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];
+        let path = ["/opt/spark/examples/src/main/resources/people.csv"];
 
         let mut df = spark
             .read()
             .format("csv")
             .option("header", "True")
             .option("delimiter", ";")
-            .load(paths);
+            .load(path);
 
         let mut df = df
             .select(col("age").cast("int").alias("age_int"))
@@ -1040,14 +1051,14 @@ mod tests {
     async fn test_drop() {
         let spark = setup().await;
 
-        let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];
+        let path = ["/opt/spark/examples/src/main/resources/people.csv"];
 
         let mut df = spark
             .read()
             .format("csv")
             .option("header", "True")
             .option("delimiter", ";")
-            .load(paths);
+            .load(path);
 
         let mut df = df.drop(vec!["age", "job"]);
 
@@ -1060,18 +1071,18 @@ mod tests {
     async fn test_join() {
         let spark = setup().await;
 
-        let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];
+        let path = ["/opt/spark/examples/src/main/resources/people.csv"];
 
         let mut df = spark
             .clone()
             .read()
             .format("csv")
             .option("header", "True")
             .option("delimiter", ";")
-            .load(paths)
+            .load(path)
             .alias("df");
 
-        let mut df1 = spark
+        let df1 = spark
             .clone()
             .range(None, 1, 1, Some(1))
             .select(vec![lit("Bob").alias("name"), lit(1).alias("id")])

src/expressions.rs

@@ -45,6 +45,15 @@ where
     }
 }
 
+impl<const N: usize, T> ToVecExpr for [T; N]
+where
+    T: ToExpr,
+{
+    fn to_vec_expr(&self) -> Vec<spark::Expression> {
+        self.iter().map(|col| col.to_expr()).collect()
+    }
+}
+
 pub trait ToFilterExpr {
     fn to_filter_expr(&self) -> Option<spark::Expression>;
 }

src/functions/mod.rs

@@ -507,14 +507,14 @@ mod tests {
     async fn test_func_col_contains() {
         let spark = setup().await;
 
-        let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];
+        let path = ["/opt/spark/examples/src/main/resources/people.csv"];
 
         let mut df = spark
             .read()
             .format("csv")
             .option("header", "True")
             .option("delimiter", ";")
-            .load(paths);
+            .load(path);
 
         let row = df
             .filter(col("name").contains("e"))
@@ -536,14 +536,14 @@ mod tests {
     async fn test_func_col_isin() {
         let spark = setup().await;
 
-        let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];
+        let path = ["/opt/spark/examples/src/main/resources/people.csv"];
 
         let mut df = spark
             .read()
             .format("csv")
             .option("header", "True")
             .option("delimiter", ";")
-            .load(paths);
+            .load(path);
 
         let row = df
             .filter(col("name").isin(vec!["Jorge", "Bob"]))

src/lib.rs

@@ -154,14 +154,14 @@ mod tests {
     async fn test_dataframe_read() {
         let spark = setup().await;
 
-        let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];
+        let path = ["/opt/spark/examples/src/main/resources/people.csv"];
 
         let mut df = spark
             .read()
             .format("csv")
             .option("header", "True")
             .option("delimiter", ";")
-            .load(paths);
+            .load(path);
 
         let rows = df
             .filter("age > 30")
@@ -197,7 +197,7 @@ mod tests {
             .read()
             .format("csv")
             .option("header", "true")
-            .load(vec![path.to_string()]);
+            .load([path]);
 
         let total: usize = df
             .select(vec![col("range_id")])