spark_connect/
merge.rs

1//! `MergeIntoWriter` mirroring `pyspark.sql.connect.merge.MergeIntoWriter`.
2//!
3//! Builds a `MergeIntoTableCommand` from a fluent set of `whenMatched` /
4//! `whenNotMatched` / `whenNotMatchedBySource` clauses and executes it.
5
6use std::collections::HashMap;
7
8use spark_connect_core::error::Result;
9use spark_connect_proto as proto;
10
11use crate::column::Column;
12use crate::dataframe::{build_input_relation, execute_command};
13use crate::plan::LogicalPlan;
14use crate::session::SparkSession;
15
16fn merge_action(
17    action_type: proto::merge_action::ActionType,
18    condition: Option<&Column>,
19    assignments: Option<HashMap<String, Column>>,
20) -> proto::Expression {
21    let mut action = proto::MergeAction::default();
22    action.action_type = action_type as i32;
23    action.condition = condition.map(|c| Box::new(c.to_proto()));
24    if let Some(assignments) = assignments {
25        action.assignments = assignments
26            .into_iter()
27            .map(|(k, v)| proto::merge_action::Assignment {
28                // Mirrors the reference `expr(k)` for the assignment target.
29                key: Some(crate::functions::expr(&k).to_proto()),
30                value: Some(v.to_proto()),
31            })
32            .collect();
33    }
34    let mut expr = proto::Expression::default();
35    expr.expr_type = Some(proto::expression::ExprType::MergeAction(Box::new(action)));
36    expr
37}
38
39/// Fluent builder for a MERGE INTO command. Create it via [`crate::dataframe::DataFrame::merge_into`].
40pub struct MergeIntoWriter {
41    session: SparkSession,
42    source_plan: LogicalPlan,
43    target_table: String,
44    condition: Column,
45    schema_evolution: bool,
46    matched_actions: Vec<proto::Expression>,
47    not_matched_actions: Vec<proto::Expression>,
48    not_matched_by_source_actions: Vec<proto::Expression>,
49}
50
51impl MergeIntoWriter {
52    pub(crate) fn new(
53        session: SparkSession,
54        source_plan: LogicalPlan,
55        target_table: String,
56        condition: Column,
57    ) -> Self {
58        MergeIntoWriter {
59            session,
60            source_plan,
61            target_table,
62            condition,
63            schema_evolution: false,
64            matched_actions: Vec::new(),
65            not_matched_actions: Vec::new(),
66            not_matched_by_source_actions: Vec::new(),
67        }
68    }
69
70    /// Begin a `WHEN MATCHED [AND condition]` clause.
71    pub fn when_matched(self, condition: Option<Column>) -> WhenMatched {
72        WhenMatched {
73            writer: self,
74            condition,
75        }
76    }
77
78    /// Begin a `WHEN NOT MATCHED [AND condition]` clause.
79    pub fn when_not_matched(self, condition: Option<Column>) -> WhenNotMatched {
80        WhenNotMatched {
81            writer: self,
82            condition,
83        }
84    }
85
86    /// Begin a `WHEN NOT MATCHED BY SOURCE [AND condition]` clause.
87    pub fn when_not_matched_by_source(self, condition: Option<Column>) -> WhenNotMatchedBySource {
88        WhenNotMatchedBySource {
89            writer: self,
90            condition,
91        }
92    }
93
94    /// Enable schema evolution for this merge.
95    pub fn with_schema_evolution(mut self) -> Self {
96        self.schema_evolution = true;
97        self
98    }
99
100    /// Execute the merge.
101    pub fn merge(self) -> Result<()> {
102        let mut cmd = proto::MergeIntoTableCommand::default();
103        cmd.target_table_name = self.target_table;
104        cmd.source_table_plan = Some(build_input_relation(&self.source_plan, &self.session)?);
105        cmd.merge_condition = Some(self.condition.to_proto());
106        cmd.match_actions = self.matched_actions;
107        cmd.not_matched_actions = self.not_matched_actions;
108        cmd.not_matched_by_source_actions = self.not_matched_by_source_actions;
109        cmd.with_schema_evolution = self.schema_evolution;
110        execute_command(
111            &self.session,
112            proto::command::CommandType::MergeIntoTableCommand(cmd),
113        )
114    }
115}
116
117/// `WHEN MATCHED` clause builder.
118pub struct WhenMatched {
119    writer: MergeIntoWriter,
120    condition: Option<Column>,
121}
122
123impl WhenMatched {
124    pub fn update_all(mut self) -> MergeIntoWriter {
125        let action = merge_action(
126            proto::merge_action::ActionType::UpdateStar,
127            self.condition.as_ref(),
128            None,
129        );
130        self.writer.matched_actions.push(action);
131        self.writer
132    }
133
134    pub fn update(mut self, assignments: HashMap<String, Column>) -> MergeIntoWriter {
135        let action = merge_action(
136            proto::merge_action::ActionType::Update,
137            self.condition.as_ref(),
138            Some(assignments),
139        );
140        self.writer.matched_actions.push(action);
141        self.writer
142    }
143
144    pub fn delete(mut self) -> MergeIntoWriter {
145        let action = merge_action(
146            proto::merge_action::ActionType::Delete,
147            self.condition.as_ref(),
148            None,
149        );
150        self.writer.matched_actions.push(action);
151        self.writer
152    }
153}
154
155/// `WHEN NOT MATCHED` clause builder.
156pub struct WhenNotMatched {
157    writer: MergeIntoWriter,
158    condition: Option<Column>,
159}
160
161impl WhenNotMatched {
162    pub fn insert_all(mut self) -> MergeIntoWriter {
163        let action = merge_action(
164            proto::merge_action::ActionType::InsertStar,
165            self.condition.as_ref(),
166            None,
167        );
168        self.writer.not_matched_actions.push(action);
169        self.writer
170    }
171
172    pub fn insert(mut self, assignments: HashMap<String, Column>) -> MergeIntoWriter {
173        let action = merge_action(
174            proto::merge_action::ActionType::Insert,
175            self.condition.as_ref(),
176            Some(assignments),
177        );
178        self.writer.not_matched_actions.push(action);
179        self.writer
180    }
181}
182
183/// `WHEN NOT MATCHED BY SOURCE` clause builder.
184pub struct WhenNotMatchedBySource {
185    writer: MergeIntoWriter,
186    condition: Option<Column>,
187}
188
189impl WhenNotMatchedBySource {
190    pub fn update_all(mut self) -> MergeIntoWriter {
191        let action = merge_action(
192            proto::merge_action::ActionType::UpdateStar,
193            self.condition.as_ref(),
194            None,
195        );
196        self.writer.not_matched_by_source_actions.push(action);
197        self.writer
198    }
199
200    pub fn update(mut self, assignments: HashMap<String, Column>) -> MergeIntoWriter {
201        let action = merge_action(
202            proto::merge_action::ActionType::Update,
203            self.condition.as_ref(),
204            Some(assignments),
205        );
206        self.writer.not_matched_by_source_actions.push(action);
207        self.writer
208    }
209
210    pub fn delete(mut self) -> MergeIntoWriter {
211        let action = merge_action(
212            proto::merge_action::ActionType::Delete,
213            self.condition.as_ref(),
214            None,
215        );
216        self.writer.not_matched_by_source_actions.push(action);
217        self.writer
218    }
219}
220
221#[cfg(test)]
222mod tests {
223    use crate::session::SparkSession;
224
225    fn session() -> SparkSession {
226        SparkSession::builder()
227            .remote("sc://localhost:15002")
228            .get_or_create()
229            .expect("failed to build session")
230    }
231
232    #[test]
233    fn merge_into_when_matched_update() {
234        let spark = session();
235        let df = spark.range(3).unwrap();
236        let mut assignments = std::collections::HashMap::new();
237        assignments.insert("col1".to_string(), crate::column::col("new_val"));
238
239        let writer = df
240            .merge_into("target_table", crate::column::col("id"))
241            .when_matched(None)
242            .update(assignments);
243
244        assert_eq!(writer.matched_actions.len(), 1);
245        assert_eq!(writer.not_matched_actions.len(), 0);
246    }
247
248    #[test]
249    fn merge_into_when_matched_delete() {
250        let spark = session();
251        let df = spark.range(3).unwrap();
252
253        let writer = df
254            .merge_into("target_table", crate::column::col("id"))
255            .when_matched(None)
256            .delete();
257
258        assert_eq!(writer.matched_actions.len(), 1);
259    }
260
261    #[test]
262    fn merge_into_when_matched_update_all() {
263        let spark = session();
264        let df = spark.range(3).unwrap();
265
266        let writer = df
267            .merge_into("target_table", crate::column::col("id"))
268            .when_matched(None)
269            .update_all();
270
271        assert_eq!(writer.matched_actions.len(), 1);
272    }
273
274    #[test]
275    fn merge_into_when_not_matched_insert() {
276        let spark = session();
277        let df = spark.range(3).unwrap();
278        let mut assignments = std::collections::HashMap::new();
279        assignments.insert("col1".to_string(), crate::column::col("val"));
280
281        let writer = df
282            .merge_into("target_table", crate::column::col("id"))
283            .when_not_matched(None)
284            .insert(assignments);
285
286        assert_eq!(writer.not_matched_actions.len(), 1);
287        assert_eq!(writer.matched_actions.len(), 0);
288    }
289
290    #[test]
291    fn merge_into_when_not_matched_insert_all() {
292        let spark = session();
293        let df = spark.range(3).unwrap();
294
295        let writer = df
296            .merge_into("target_table", crate::column::col("id"))
297            .when_not_matched(None)
298            .insert_all();
299
300        assert_eq!(writer.not_matched_actions.len(), 1);
301    }
302
303    #[test]
304    fn merge_into_when_not_matched_by_source() {
305        let spark = session();
306        let df = spark.range(3).unwrap();
307
308        let writer = df
309            .merge_into("target_table", crate::column::col("id"))
310            .when_not_matched_by_source(None)
311            .delete();
312
313        assert_eq!(writer.not_matched_by_source_actions.len(), 1);
314    }
315
316    #[test]
317    fn merge_into_multiple_clauses() {
318        let spark = session();
319        let df = spark.range(3).unwrap();
320        let mut assign1 = std::collections::HashMap::new();
321        assign1.insert("col1".to_string(), crate::column::col("val1"));
322        let mut assign2 = std::collections::HashMap::new();
323        assign2.insert("col2".to_string(), crate::column::col("val2"));
324
325        let writer = df
326            .merge_into("target_table", crate::column::col("id"))
327            .when_matched(None)
328            .update(assign1)
329            .when_not_matched(None)
330            .insert(assign2);
331
332        assert_eq!(writer.matched_actions.len(), 1);
333        assert_eq!(writer.not_matched_actions.len(), 1);
334    }
335
336    #[test]
337    fn merge_into_with_schema_evolution() {
338        let spark = session();
339        let df = spark.range(3).unwrap();
340
341        let writer = df
342            .merge_into("target_table", crate::column::col("id"))
343            .with_schema_evolution()
344            .when_matched(None)
345            .update_all();
346
347        assert!(writer.schema_evolution);
348    }
349}
spark_connect/merge.rs

spark_connect/
merge.rs
