spark_connect/
readwriter.rs

1//! DataFrameReader and DataFrameWriter implementation mirroring `pyspark.sql.connect.readwriter`.
2//!
3//! Provides the API for reading data from various sources and writing DataFrames to files/tables.
4
5use std::collections::HashMap;
6
7use spark_connect_core::error::Result;
8use spark_connect_proto as proto;
9
10use crate::column::Column;
11use crate::dataframe::{build_input_relation, execute_command, DataFrame};
12use crate::plan::LogicalPlan;
13use crate::session::SparkSession;
14
15/// DataFrameReader for reading data from various sources.
16///
17/// Mirrors `pyspark.sql.connect.readwriter.DataFrameReader`.
18pub struct DataFrameReader {
19    session: SparkSession,
20    format: Option<String>,
21    schema: String,
22    options: HashMap<String, String>,
23}
24
25impl DataFrameReader {
26    /// Create a new DataFrameReader.
27    pub(crate) fn new(session: SparkSession) -> Self {
28        DataFrameReader {
29            session,
30            format: None,
31            schema: String::new(),
32            options: HashMap::new(),
33        }
34    }
35
36    /// Set the format/source type (e.g., "json", "parquet", "csv").
37    pub fn format(mut self, source: &str) -> Self {
38        self.format = Some(source.to_string());
39        self
40    }
41
42    /// Set the schema from a DDL string or JSON string.
43    pub fn schema(mut self, schema: String) -> Self {
44        self.schema = schema;
45        self
46    }
47
48    /// Set a single option key-value pair.
49    pub fn option(mut self, key: &str, value: &str) -> Self {
50        self.options.insert(key.to_string(), value.to_string());
51        self
52    }
53
54    /// Set multiple options.
55    pub fn options(mut self, options: HashMap<String, String>) -> Self {
56        self.options.extend(options);
57        self
58    }
59
60    /// Load data from the specified path(s) with the configured format, schema, and options.
61    pub fn load(self, path: Option<&str>) -> DataFrame {
62        let paths = path.map(|p| vec![p.to_string()]);
63        let plan = LogicalPlan::Read {
64            read_type: ReadType::DataSource {
65                format: self.format.clone(),
66                schema: if self.schema.is_empty() {
67                    None
68                } else {
69                    Some(self.schema.clone())
70                },
71                options: self.options.clone(),
72                paths: paths.unwrap_or_default(),
73                predicates: vec![],
74                source_name: None,
75            },
76            is_streaming: false,
77        };
78        DataFrame::new(self.session, plan)
79    }
80
81    /// Read the CDC changes of a named table. Mirrors `DataFrameReader.changes`.
82    pub fn changes(self, table_name: &str) -> DataFrame {
83        let plan = LogicalPlan::RelationChanges {
84            table_name: table_name.to_string(),
85            options: self.options.clone(),
86            is_streaming: None,
87        };
88        DataFrame::new(self.session, plan)
89    }
90
91    /// Read from a named table.
92    pub fn table(self, table_name: &str) -> DataFrame {
93        let plan = LogicalPlan::Read {
94            read_type: ReadType::NamedTable {
95                table_name: table_name.to_string(),
96                options: self.options.clone(),
97            },
98            is_streaming: false,
99        };
100        DataFrame::new(self.session, plan)
101    }
102
103    /// Read JSON data.
104    pub fn json(mut self, path: &str) -> DataFrame {
105        self.format = Some("json".to_string());
106        let paths = vec![path.to_string()];
107        let plan = LogicalPlan::Read {
108            read_type: ReadType::DataSource {
109                format: self.format.clone(),
110                schema: if self.schema.is_empty() {
111                    None
112                } else {
113                    Some(self.schema.clone())
114                },
115                options: self.options.clone(),
116                paths,
117                predicates: vec![],
118                source_name: None,
119            },
120            is_streaming: false,
121        };
122        DataFrame::new(self.session, plan)
123    }
124
125    /// Read Parquet data.
126    pub fn parquet(mut self, path: &str) -> DataFrame {
127        self.format = Some("parquet".to_string());
128        let paths = vec![path.to_string()];
129        let plan = LogicalPlan::Read {
130            read_type: ReadType::DataSource {
131                format: self.format.clone(),
132                schema: if self.schema.is_empty() {
133                    None
134                } else {
135                    Some(self.schema.clone())
136                },
137                options: self.options.clone(),
138                paths,
139                predicates: vec![],
140                source_name: None,
141            },
142            is_streaming: false,
143        };
144        DataFrame::new(self.session, plan)
145    }
146
147    /// Read CSV data.
148    pub fn csv(mut self, path: &str) -> DataFrame {
149        self.format = Some("csv".to_string());
150        let paths = vec![path.to_string()];
151        let plan = LogicalPlan::Read {
152            read_type: ReadType::DataSource {
153                format: self.format.clone(),
154                schema: if self.schema.is_empty() {
155                    None
156                } else {
157                    Some(self.schema.clone())
158                },
159                options: self.options.clone(),
160                paths,
161                predicates: vec![],
162                source_name: None,
163            },
164            is_streaming: false,
165        };
166        DataFrame::new(self.session, plan)
167    }
168
169    /// Read ORC data.
170    pub fn orc(mut self, path: &str) -> DataFrame {
171        self.format = Some("orc".to_string());
172        let paths = vec![path.to_string()];
173        let plan = LogicalPlan::Read {
174            read_type: ReadType::DataSource {
175                format: self.format.clone(),
176                schema: if self.schema.is_empty() {
177                    None
178                } else {
179                    Some(self.schema.clone())
180                },
181                options: self.options.clone(),
182                paths,
183                predicates: vec![],
184                source_name: None,
185            },
186            is_streaming: false,
187        };
188        DataFrame::new(self.session, plan)
189    }
190
191    /// Read text data.
192    pub fn text(mut self, path: &str) -> DataFrame {
193        self.format = Some("text".to_string());
194        let paths = vec![path.to_string()];
195        let plan = LogicalPlan::Read {
196            read_type: ReadType::DataSource {
197                format: self.format.clone(),
198                schema: if self.schema.is_empty() {
199                    None
200                } else {
201                    Some(self.schema.clone())
202                },
203                options: self.options.clone(),
204                paths,
205                predicates: vec![],
206                source_name: None,
207            },
208            is_streaming: false,
209        };
210        DataFrame::new(self.session, plan)
211    }
212
213    /// Read XML file(s). Mirrors `DataFrameReader.xml`.
214    pub fn xml(mut self, path: &str) -> DataFrame {
215        self.format = Some("xml".to_string());
216        let paths = vec![path.to_string()];
217        let plan = LogicalPlan::Read {
218            read_type: ReadType::DataSource {
219                format: self.format.clone(),
220                schema: if self.schema.is_empty() {
221                    None
222                } else {
223                    Some(self.schema.clone())
224                },
225                options: self.options.clone(),
226                paths,
227                predicates: vec![],
228                source_name: None,
229            },
230            is_streaming: false,
231        };
232        DataFrame::new(self.session, plan)
233    }
234
235    /// Read from JDBC data source.
236    pub fn jdbc(mut self, url: &str, table: &str, predicates: Option<Vec<String>>) -> DataFrame {
237        self.format = Some("jdbc".to_string());
238        self.options.insert("url".to_string(), url.to_string());
239        self.options
240            .insert("dbtable".to_string(), table.to_string());
241
242        let plan = LogicalPlan::Read {
243            read_type: ReadType::DataSource {
244                format: self.format.clone(),
245                schema: if self.schema.is_empty() {
246                    None
247                } else {
248                    Some(self.schema.clone())
249                },
250                options: self.options.clone(),
251                paths: vec![],
252                predicates: predicates.unwrap_or_default(),
253                source_name: None,
254            },
255            is_streaming: false,
256        };
257        DataFrame::new(self.session, plan)
258    }
259}
260
261/// ReadType variant for Read relation.
262#[derive(Debug, Clone)]
263pub enum ReadType {
264    /// DataSource read (file-based or other formats).
265    DataSource {
266        format: Option<String>,
267        schema: Option<String>,
268        options: HashMap<String, String>,
269        paths: Vec<String>,
270        predicates: Vec<String>,
271        source_name: Option<String>,
272    },
273    /// Named table read.
274    NamedTable {
275        table_name: String,
276        options: HashMap<String, String>,
277    },
278}
279
280/// SaveMode for write operations.
281#[derive(Debug, Clone, Copy, PartialEq, Eq)]
282pub enum SaveMode {
283    Append,
284    Overwrite,
285    ErrorIfExists,
286    Ignore,
287}
288
289impl SaveMode {
290    /// Convert SaveMode to proto i32 value.
291    pub fn to_proto(&self) -> i32 {
292        match self {
293            SaveMode::Append => 1i32,
294            SaveMode::Overwrite => 2i32,
295            SaveMode::ErrorIfExists => 3i32,
296            SaveMode::Ignore => 4i32,
297        }
298    }
299
300    /// Parse SaveMode from string.
301    pub fn from_str(s: &str) -> Option<Self> {
302        match s.to_lowercase().as_str() {
303            "append" => Some(SaveMode::Append),
304            "overwrite" => Some(SaveMode::Overwrite),
305            "error" | "errorifexists" => Some(SaveMode::ErrorIfExists),
306            "ignore" => Some(SaveMode::Ignore),
307            _ => None,
308        }
309    }
310}
311
312/// TableSaveMethod for WriteOperation.
313#[derive(Debug, Clone, Copy, PartialEq, Eq)]
314pub enum TableSaveMethod {
315    SaveAsTable,
316    InsertInto,
317}
318
319impl TableSaveMethod {
320    /// Convert to proto i32 value.
321    pub fn to_proto(&self) -> i32 {
322        match self {
323            TableSaveMethod::SaveAsTable => 1i32,
324            TableSaveMethod::InsertInto => 2i32,
325        }
326    }
327}
328
329/// DataFrameWriter for writing DataFrames to various destinations.
330///
331/// Mirrors `pyspark.sql.connect.readwriter.DataFrameWriter`.
332pub struct DataFrameWriter {
333    session: SparkSession,
334    input_plan: LogicalPlan,
335    format: Option<String>,
336    mode: SaveMode,
337    options: HashMap<String, String>,
338    partition_cols: Vec<String>,
339    cluster_cols: Vec<String>,
340    bucket_cols: Vec<String>,
341    sort_cols: Vec<String>,
342    num_buckets: Option<i32>,
343}
344
345impl DataFrameWriter {
346    /// Create a new DataFrameWriter.
347    pub(crate) fn new(session: SparkSession, input_plan: LogicalPlan) -> Self {
348        DataFrameWriter {
349            session,
350            input_plan,
351            format: None,
352            mode: SaveMode::ErrorIfExists,
353            options: HashMap::new(),
354            partition_cols: vec![],
355            cluster_cols: vec![],
356            bucket_cols: vec![],
357            sort_cols: vec![],
358            num_buckets: None,
359        }
360    }
361
362    /// Cluster the output by the given columns (liquid clustering).
363    pub fn cluster_by<S: Into<String>>(mut self, cols: impl IntoIterator<Item = S>) -> Self {
364        self.cluster_cols = cols.into_iter().map(Into::into).collect();
365        self
366    }
367
368    /// Set the save mode.
369    pub fn mode(mut self, mode: &str) -> Self {
370        if let Some(m) = SaveMode::from_str(mode) {
371            self.mode = m;
372        }
373        self
374    }
375
376    /// Set the format/source type.
377    pub fn format(mut self, source: &str) -> Self {
378        self.format = Some(source.to_string());
379        self
380    }
381
382    /// Set a single option.
383    pub fn option(mut self, key: &str, value: &str) -> Self {
384        self.options.insert(key.to_string(), value.to_string());
385        self
386    }
387
388    /// Set multiple options.
389    pub fn options(mut self, options: HashMap<String, String>) -> Self {
390        self.options.extend(options);
391        self
392    }
393
394    /// Set the partition columns.
395    pub fn partition_by<S: Into<String>>(mut self, cols: impl IntoIterator<Item = S>) -> Self {
396        self.partition_cols = cols.into_iter().map(Into::into).collect();
397        self
398    }
399
400    /// Set the bucketing specification.
401    pub fn bucket_by<S: Into<String>>(
402        mut self,
403        num_buckets: i32,
404        cols: impl IntoIterator<Item = S>,
405    ) -> Self {
406        self.num_buckets = Some(num_buckets);
407        self.bucket_cols = cols.into_iter().map(Into::into).collect();
408        self
409    }
410
411    /// Set the sort columns.
412    pub fn sort_by<S: Into<String>>(mut self, cols: impl IntoIterator<Item = S>) -> Self {
413        self.sort_cols = cols.into_iter().map(Into::into).collect();
414        self
415    }
416
417    /// Build the `WriteOperation` command proto for this writer.
418    pub(crate) fn build_write_operation(
419        &self,
420        save_type: Option<proto::write_operation::SaveType>,
421    ) -> Result<proto::WriteOperation> {
422        let mut op = proto::WriteOperation::default();
423        op.input = Some(build_input_relation(&self.input_plan, &self.session)?);
424        op.source = self.format.clone();
425        op.mode = self.mode.to_proto();
426        op.sort_column_names = self.sort_cols.clone();
427        op.partitioning_columns = self.partition_cols.clone();
428        op.clustering_columns = self.cluster_cols.clone();
429        op.options = self.options.clone();
430        op.save_type = save_type;
431        if let Some(num_buckets) = self.num_buckets {
432            let mut bucket_by = proto::write_operation::BucketBy::default();
433            bucket_by.num_buckets = num_buckets;
434            bucket_by.bucket_column_names = self.bucket_cols.clone();
435            op.bucket_by = Some(bucket_by);
436        }
437        Ok(op)
438    }
439
440    fn save_table(self, table_name: &str, method: TableSaveMethod) -> Result<()> {
441        let mut table = proto::write_operation::SaveTable::default();
442        table.table_name = table_name.to_string();
443        table.save_method = method.to_proto();
444        let op =
445            self.build_write_operation(Some(proto::write_operation::SaveType::Table(table)))?;
446        execute_command(
447            &self.session,
448            proto::command::CommandType::WriteOperation(op),
449        )
450    }
451
452    /// Save the DataFrame to a file path (or to a path-less sink such as `noop`).
453    pub fn save(self, path: Option<&str>) -> Result<()> {
454        let save_type = path.map(|p| proto::write_operation::SaveType::Path(p.to_string()));
455        let op = self.build_write_operation(save_type)?;
456        execute_command(
457            &self.session,
458            proto::command::CommandType::WriteOperation(op),
459        )
460    }
461
462    /// Save the DataFrame as a managed table.
463    pub fn save_as_table(self, table_name: &str) -> Result<()> {
464        self.save_table(table_name, TableSaveMethod::SaveAsTable)
465    }
466
467    /// Insert the DataFrame into an existing table.
468    pub fn insert_into(self, table_name: &str) -> Result<()> {
469        self.save_table(table_name, TableSaveMethod::InsertInto)
470    }
471
472    /// Write as JSON.
473    pub fn json(mut self, path: &str) -> Result<()> {
474        self.format = Some("json".to_string());
475        self.save(Some(path))
476    }
477
478    /// Write as Parquet.
479    pub fn parquet(mut self, path: &str) -> Result<()> {
480        self.format = Some("parquet".to_string());
481        self.save(Some(path))
482    }
483
484    /// Write as CSV.
485    pub fn csv(mut self, path: &str) -> Result<()> {
486        self.format = Some("csv".to_string());
487        self.save(Some(path))
488    }
489
490    /// Write as ORC.
491    pub fn orc(mut self, path: &str) -> Result<()> {
492        self.format = Some("orc".to_string());
493        self.save(Some(path))
494    }
495
496    /// Write as text.
497    pub fn text(mut self, path: &str) -> Result<()> {
498        self.format = Some("text".to_string());
499        self.save(Some(path))
500    }
501
502    /// Write as XML. Mirrors `DataFrameWriter.xml`.
503    pub fn xml(mut self, path: &str) -> Result<()> {
504        self.format = Some("xml".to_string());
505        self.save(Some(path))
506    }
507}
508
509/// DataFrameWriterV2 for the v2 write API (`DataFrame.writeTo`).
510///
511/// Mirrors `pyspark.sql.connect.readwriter.DataFrameWriterV2`.
512pub struct DataFrameWriterV2 {
513    session: SparkSession,
514    input_plan: LogicalPlan,
515    table_name: String,
516    provider: Option<String>,
517    options: HashMap<String, String>,
518    table_properties: HashMap<String, String>,
519    partition_cols: Vec<Column>,
520    cluster_cols: Vec<String>,
521}
522
523impl DataFrameWriterV2 {
524    /// Create a new DataFrameWriterV2 targeting `table_name`.
525    pub(crate) fn new(session: SparkSession, input_plan: LogicalPlan, table_name: &str) -> Self {
526        DataFrameWriterV2 {
527            session,
528            input_plan,
529            table_name: table_name.to_string(),
530            provider: None,
531            options: HashMap::new(),
532            table_properties: HashMap::new(),
533            partition_cols: vec![],
534            cluster_cols: vec![],
535        }
536    }
537
538    /// Cluster the output table by the given columns (liquid clustering).
539    pub fn cluster_by<S: Into<String>>(mut self, cols: impl IntoIterator<Item = S>) -> Self {
540        self.cluster_cols = cols.into_iter().map(Into::into).collect();
541        self
542    }
543
544    /// Specify the underlying output data source provider (e.g. "parquet").
545    pub fn using(mut self, provider: &str) -> Self {
546        self.provider = Some(provider.to_string());
547        self
548    }
549
550    /// Add a write option.
551    pub fn option(mut self, key: &str, value: &str) -> Self {
552        self.options.insert(key.to_string(), value.to_string());
553        self
554    }
555
556    /// Add multiple write options.
557    pub fn options(mut self, options: HashMap<String, String>) -> Self {
558        self.options.extend(options);
559        self
560    }
561
562    /// Add a table property.
563    pub fn table_property(mut self, property: &str, value: &str) -> Self {
564        self.table_properties
565            .insert(property.to_string(), value.to_string());
566        self
567    }
568
569    /// Partition the output table by the given columns.
570    pub fn partition_by<C: Into<Column>>(mut self, columns: impl IntoIterator<Item = C>) -> Self {
571        self.partition_cols = columns.into_iter().map(Into::into).collect();
572        self
573    }
574
575    /// Build the `WriteOperationV2` command proto for the given mode.
576    pub(crate) fn build_operation(
577        &self,
578        mode: proto::write_operation_v2::Mode,
579        overwrite_condition: Option<proto::Expression>,
580    ) -> Result<proto::WriteOperationV2> {
581        let mut op = proto::WriteOperationV2::default();
582        op.input = Some(build_input_relation(&self.input_plan, &self.session)?);
583        op.table_name = self.table_name.clone();
584        op.provider = self.provider.clone();
585        op.partitioning_columns = self.partition_cols.iter().map(|c| c.to_proto()).collect();
586        op.clustering_columns = self.cluster_cols.clone();
587        op.options = self.options.clone();
588        op.table_properties = self.table_properties.clone();
589        op.mode = mode as i32;
590        op.overwrite_condition = overwrite_condition;
591        Ok(op)
592    }
593
594    fn execute(self, mode: proto::write_operation_v2::Mode) -> Result<()> {
595        let op = self.build_operation(mode, None)?;
596        execute_command(
597            &self.session,
598            proto::command::CommandType::WriteOperationV2(op),
599        )
600    }
601
602    /// Create a new table from the DataFrame.
603    pub fn create(self) -> Result<()> {
604        self.execute(proto::write_operation_v2::Mode::Create)
605    }
606
607    /// Replace an existing table with the DataFrame.
608    pub fn replace(self) -> Result<()> {
609        self.execute(proto::write_operation_v2::Mode::Replace)
610    }
611
612    /// Create the table, or replace it if it already exists.
613    pub fn create_or_replace(self) -> Result<()> {
614        self.execute(proto::write_operation_v2::Mode::CreateOrReplace)
615    }
616
617    /// Append the DataFrame's rows to the table.
618    pub fn append(self) -> Result<()> {
619        self.execute(proto::write_operation_v2::Mode::Append)
620    }
621
622    /// Overwrite rows matching `condition` with the DataFrame's rows.
623    pub fn overwrite(self, condition: Column) -> Result<()> {
624        let op = self.build_operation(
625            proto::write_operation_v2::Mode::Overwrite,
626            Some(condition.to_proto()),
627        )?;
628        execute_command(
629            &self.session,
630            proto::command::CommandType::WriteOperationV2(op),
631        )
632    }
633
634    /// Overwrite all partitions touched by the DataFrame (dynamic overwrite).
635    pub fn overwrite_partitions(self) -> Result<()> {
636        self.execute(proto::write_operation_v2::Mode::OverwritePartitions)
637    }
638}
639
640#[cfg(test)]
641mod tests {
642    use super::*;
643    use crate::session::SparkSession;
644
645    // The gRPC channel connects lazily, so a session can be built offline for
646    // tests that only construct request protos (no server round-trip).
647    fn session() -> SparkSession {
648        SparkSession::builder()
649            .remote("sc://localhost:15002")
650            .get_or_create()
651            .expect("failed to build session")
652    }
653
654    #[test]
655    fn v1_write_operation_to_path() {
656        let spark = session();
657        let df = spark.range(3).unwrap();
658        let op = df
659            .write()
660            .format("parquet")
661            .mode("overwrite")
662            .option("compression", "snappy")
663            .partition_by(vec!["a".to_string()])
664            .build_write_operation(Some(proto::write_operation::SaveType::Path(
665                "/tmp/out".to_string(),
666            )))
667            .unwrap();
668
669        assert!(op.input.is_some());
670        assert_eq!(op.source.as_deref(), Some("parquet"));
671        assert_eq!(op.mode, SaveMode::Overwrite.to_proto());
672        assert_eq!(
673            op.options.get("compression").map(String::as_str),
674            Some("snappy")
675        );
676        assert_eq!(op.partitioning_columns, vec!["a".to_string()]);
677        match op.save_type {
678            Some(proto::write_operation::SaveType::Path(p)) => assert_eq!(p, "/tmp/out"),
679            other => panic!("expected Path save_type, got {other:?}"),
680        }
681    }
682
683    #[test]
684    fn v1_write_operation_save_as_table() {
685        let spark = session();
686        let df = spark.range(3).unwrap();
687        let mut table = proto::write_operation::SaveTable::default();
688        table.table_name = "db.people".to_string();
689        table.save_method = TableSaveMethod::SaveAsTable.to_proto();
690        let op = df
691            .write()
692            .build_write_operation(Some(proto::write_operation::SaveType::Table(table)))
693            .unwrap();
694
695        match op.save_type {
696            Some(proto::write_operation::SaveType::Table(t)) => {
697                assert_eq!(t.table_name, "db.people");
698                assert_eq!(t.save_method, TableSaveMethod::SaveAsTable.to_proto());
699            }
700            other => panic!("expected Table save_type, got {other:?}"),
701        }
702    }
703
704    #[test]
705    fn v2_write_operation_fields_and_modes() {
706        let spark = session();
707        let df = spark.range(3).unwrap();
708        let op = df
709            .write_to("db.tbl")
710            .using("delta")
711            .option("mergeSchema", "true")
712            .table_property("owner", "eng")
713            .partition_by(vec![crate::column::col("a")])
714            .build_operation(proto::write_operation_v2::Mode::Append, None)
715            .unwrap();
716
717        assert!(op.input.is_some());
718        assert_eq!(op.table_name, "db.tbl");
719        assert_eq!(op.provider.as_deref(), Some("delta"));
720        assert_eq!(op.mode, proto::write_operation_v2::Mode::Append as i32);
721        assert_eq!(op.partitioning_columns.len(), 1);
722        assert_eq!(
723            op.options.get("mergeSchema").map(String::as_str),
724            Some("true")
725        );
726        assert_eq!(
727            op.table_properties.get("owner").map(String::as_str),
728            Some("eng")
729        );
730
731        // Terminal modes map to the correct proto enum values.
732        let create = df
733            .write_to("t")
734            .build_operation(proto::write_operation_v2::Mode::Create, None)
735            .unwrap();
736        assert_eq!(create.mode, 1);
737        let cor = df
738            .write_to("t")
739            .build_operation(proto::write_operation_v2::Mode::CreateOrReplace, None)
740            .unwrap();
741        assert_eq!(cor.mode, 6);
742    }
743
744    #[test]
745    fn v2_overwrite_sets_condition() {
746        let spark = session();
747        let df = spark.range(3).unwrap();
748        let op = df
749            .write_to("t")
750            .build_operation(
751                proto::write_operation_v2::Mode::Overwrite,
752                Some(crate::column::col("id").to_proto()),
753            )
754            .unwrap();
755        assert_eq!(op.mode, proto::write_operation_v2::Mode::Overwrite as i32);
756        assert!(op.overwrite_condition.is_some());
757    }
758
759    #[test]
760    fn v1_write_operation_cluster_by() {
761        let spark = session();
762        let df = spark.range(3).unwrap();
763        let op = df
764            .write()
765            .format("parquet")
766            .cluster_by(vec!["col1".to_string(), "col2".to_string()])
767            .build_write_operation(Some(proto::write_operation::SaveType::Path(
768                "/tmp/out".to_string(),
769            )))
770            .unwrap();
771
772        assert_eq!(op.clustering_columns, vec!["col1", "col2"]);
773    }
774
775    #[test]
776    fn v2_write_operation_cluster_by() {
777        let spark = session();
778        let df = spark.range(3).unwrap();
779        let op = df
780            .write_to("t")
781            .cluster_by(vec!["col1".to_string(), "col2".to_string()])
782            .build_operation(proto::write_operation_v2::Mode::Create, None)
783            .unwrap();
784
785        assert_eq!(op.clustering_columns, vec!["col1", "col2"]);
786    }
787
788    #[test]
789    fn reader_jdbc_with_options() {
790        let spark = session();
791        let reader = spark.read();
792        let df = reader
793            .option("url", "jdbc:mysql://localhost:3306/db")
794            .option("user", "root")
795            .option("password", "secret")
796            .jdbc("jdbc:mysql://localhost:3306/db", "table_name", None);
797
798        match &df.plan {
799            LogicalPlan::Read {
800                read_type:
801                    ReadType::DataSource {
802                        options, format, ..
803                    },
804                ..
805            } => {
806                assert_eq!(format.as_deref(), Some("jdbc"));
807                assert_eq!(
808                    options.get("url").map(String::as_str),
809                    Some("jdbc:mysql://localhost:3306/db")
810                );
811                assert_eq!(options.get("user").map(String::as_str), Some("root"));
812                assert_eq!(options.get("password").map(String::as_str), Some("secret"));
813            }
814            _ => panic!("expected Read plan"),
815        }
816    }
817
818    #[test]
819    fn reader_jdbc_with_predicates() {
820        let spark = session();
821        let reader = spark.read();
822        let predicates = vec!["col1 > 10".to_string(), "col2 = 'value'".to_string()];
823        let df = reader.jdbc(
824            "jdbc:mysql://localhost/db",
825            "table",
826            Some(predicates.clone()),
827        );
828
829        match &df.plan {
830            LogicalPlan::Read {
831                read_type:
832                    ReadType::DataSource {
833                        predicates: preds, ..
834                    },
835                ..
836            } => {
837                assert_eq!(preds.len(), 2);
838            }
839            _ => panic!("expected Read plan with predicates"),
840        }
841    }
842
843    #[test]
844    fn v1_write_partition_and_cluster() {
845        let spark = session();
846        let df = spark.range(3).unwrap();
847        let op = df
848            .write()
849            .format("delta")
850            .partition_by(vec!["date".to_string()])
851            .cluster_by(vec!["user_id".to_string()])
852            .build_write_operation(Some(proto::write_operation::SaveType::Path(
853                "/tmp/data".to_string(),
854            )))
855            .unwrap();
856
857        assert_eq!(op.partitioning_columns, vec!["date"]);
858        assert_eq!(op.clustering_columns, vec!["user_id"]);
859    }
860
861    #[test]
862    fn v1_write_bucket_by() {
863        let spark = session();
864        let df = spark.range(3).unwrap();
865        let op = df
866            .write()
867            .format("parquet")
868            .bucket_by(10, vec!["col1".to_string()])
869            .build_write_operation(Some(proto::write_operation::SaveType::Path(
870                "/tmp/out".to_string(),
871            )))
872            .unwrap();
873
874        assert!(op.bucket_by.is_some());
875        let bucket_by = op.bucket_by.unwrap();
876        assert_eq!(bucket_by.num_buckets, 10);
877        assert_eq!(bucket_by.bucket_column_names, vec!["col1"]);
878    }
879
880    #[test]
881    fn v1_write_sort_by() {
882        let spark = session();
883        let df = spark.range(3).unwrap();
884        let op = df
885            .write()
886            .format("parquet")
887            .sort_by(vec!["col1".to_string()])
888            .build_write_operation(Some(proto::write_operation::SaveType::Path(
889                "/tmp/out".to_string(),
890            )))
891            .unwrap();
892
893        assert_eq!(op.sort_column_names, vec!["col1"]);
894    }
895
896    #[test]
897    fn reader_format_options() {
898        let spark = session();
899        let mut opts = std::collections::HashMap::new();
900        opts.insert("delimiter".to_string(), ";".to_string());
901        opts.insert("header".to_string(), "true".to_string());
902
903        let df = spark
904            .read()
905            .format("csv")
906            .options(opts)
907            .load(Some("/data.csv"));
908
909        match &df.plan {
910            LogicalPlan::Read {
911                read_type:
912                    ReadType::DataSource {
913                        options, format, ..
914                    },
915                ..
916            } => {
917                assert_eq!(format.as_deref(), Some("csv"));
918                assert_eq!(options.get("delimiter").map(String::as_str), Some(";"));
919                assert_eq!(options.get("header").map(String::as_str), Some("true"));
920            }
921            _ => panic!("expected Read plan"),
922        }
923    }
924}
spark_connect/readwriter.rs

spark_connect/
readwriter.rs
