1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

//! Spark Connection Client for Rust
//!
//! Currently, the Spark Connect client for Rust is **highly experimental** and **should
//! not be used in any production setting**. This is currently a "proof of concept" to identify the methods
//! of interacting with Spark cluster from rust.
//!
//! # Usage
//!
//! Create a Spark Session and create a DataFrame from a SQL statement:
//!
//! ```rust
//! async {
//!
//!     let spark: SparkSession = SparkSessionBuilder::remote("sc://127.0.0.1:15002/;user_id=example_rs".to_string())
//!         .build()
//!         .await?;
//!
//!     let mut df = spark.sql("SELECT * FROM json.`/opt/spark/examples/src/main/resources/employees.json`");
//!
//!     df.filter("salary > 3000").show(Some(5), None, None).await?;
//! };
//!```
//!
//! Create a Spark Session, create a DataFrame from a CSV file, and write the results:
//!
//! ```rust
//! async {
//!
//!     let spark: SparkSession = SparkSessionBuilder::remote("sc://127.0.0.1:15002/;user_id=example_rs".to_string())
//!         .build()
//!         .await?;
//!
//!     let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];
//!
//!     let mut df = spark
//!         .read()
//!         .format("csv")
//!         .option("header", "True")
//!         .option("delimiter", ";")
//!         .load(paths);
//!
//!     let mut df = df
//!         .filter("age > 30")
//!         .select(vec!["name"]);
//!
//!     df.write()
//!       .format("csv")
//!       .option("header", "true")
//!       .save("/opt/spark/examples/src/main/rust/people/")
//!       .await?;
//! };
//!```
//!

/// Spark Connect gRPC protobuf translated using [tonic]
pub mod spark {
    tonic::include_proto!("spark.connect");
}

pub mod dataframe;
pub mod execution;
pub mod plan;

pub use arrow;
pub use dataframe::{DataFrame, DataFrameReader, DataFrameWriter};
pub use execution::context::{SparkSession, SparkSessionBuilder};
pub use plan::LogicalPlanBuilder;

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

    async fn setup() -> SparkSession {
        println!("SparkSession Setup");

        let connection = "sc://127.0.0.1:15002/;user_id=rust_test".to_string();

        SparkSessionBuilder::remote(connection)
            .build()
            .await
            .unwrap()
    }

    #[tokio::test]
    async fn test_dataframe_range() {
        let spark = setup().await;

        let mut df = spark.range(None, 100, 1, Some(8));

        let rows = df.collect().await.unwrap();

        let total: usize = rows.iter().map(|batch| batch.num_rows()).sum();

        assert_eq!(total, 100)
    }

    #[tokio::test]
    async fn test_dataframe_read() {
        let spark = setup().await;

        let paths = vec!["/opt/spark/examples/src/main/resources/people.csv".to_string()];

        let mut df = spark
            .read()
            .format("csv")
            .option("header", "True")
            .option("delimiter", ";")
            .load(paths);

        let rows = df
            .filter("age > 30")
            .select(vec!["name"])
            .collect()
            .await
            .unwrap();

        assert_eq!(rows[0].num_rows(), 1);
        // assert_eq!(rows[0].column(0).);
    }

    #[tokio::test]
    async fn test_dataframe_write() {
        let spark = setup().await;

        let df = spark
            .clone()
            .range(None, 1000, 1, Some(16))
            .selectExpr(vec!["id AS range_id"]);

        let path = "/opt/spark/examples/src/main/rust/employees/";

        df.write()
            .format("csv")
            .option("header", "true")
            .save(path)
            .await
            .unwrap();

        let mut df = spark
            .clone()
            .read()
            .format("csv")
            .option("header", "true")
            .load(vec![path.to_string()]);

        let total: usize = df
            .select(vec!["range_id"])
            .collect()
            .await
            .unwrap()
            .iter()
            .map(|batch| batch.num_rows())
            .sum();

        assert_eq!(total, 1000)
    }
}