@@ -23,7 +23,7 @@ use std::sync::Arc;
2323
2424use arrow:: datatypes:: { DataType , Field , Schema } ;
2525use datafusion_common:: file_options:: file_type:: FileType ;
26- use datafusion_common:: { DFSchemaRef , TableReference } ;
26+ use datafusion_common:: { DFSchemaRef , Result , TableReference , internal_err } ;
2727
2828use crate :: { Expr , LogicalPlan , TableSource } ;
2929
@@ -299,12 +299,129 @@ impl Display for InsertOp {
299299}
300300
301301/// Describes a MERGE INTO operation's parameters.
302- #[ derive( Debug , Clone , PartialEq , Eq , PartialOrd , Hash ) ]
302+ #[ derive( Debug , Clone , PartialEq , Eq , Hash ) ]
303303pub struct MergeIntoOp {
304304 /// The join condition from `ON <expr>`.
305305 pub on : Expr ,
306306 /// The WHEN clauses, in the order they appeared in the SQL.
307307 pub clauses : Vec < MergeIntoClause > ,
308+ /// Schema of the target table (qualified with the alias or table name used
309+ /// in the SQL). Stored here so the analyzer can build a combined
310+ /// target+source schema for type coercion and function rewrites.
311+ pub target_schema : DFSchemaRef ,
312+ }
313+
314+ impl PartialOrd for MergeIntoOp {
315+ fn partial_cmp ( & self , other : & Self ) -> Option < Ordering > {
316+ match self . on . partial_cmp ( & other. on ) {
317+ Some ( Ordering :: Equal ) => self . clauses . partial_cmp ( & other. clauses ) ,
318+ cmp => cmp,
319+ }
320+ }
321+ }
322+
323+ impl MergeIntoOp {
324+ /// Count of top-level [`Expr`]s owned by this operation (no allocation).
325+ ///
326+ /// Matches the length of [`Self::exprs`] and the `exprs` vec consumed by
327+ /// [`Self::with_new_exprs`].
328+ fn expr_count ( & self ) -> usize {
329+ 1 + self
330+ . clauses
331+ . iter ( )
332+ . map ( |c| {
333+ c. predicate . is_some ( ) as usize
334+ + match & c. action {
335+ MergeIntoAction :: Update ( a) => a. len ( ) ,
336+ MergeIntoAction :: Insert { values, .. } => values. len ( ) ,
337+ MergeIntoAction :: Delete => 0 ,
338+ }
339+ } )
340+ . sum :: < usize > ( )
341+ }
342+
343+ /// Top-level [`Expr`]s in stable order: `on`, then per-clause predicate
344+ /// (if any) and action value expressions.
345+ pub fn exprs ( & self ) -> Vec < & Expr > {
346+ let mut out = Vec :: with_capacity ( self . expr_count ( ) ) ;
347+ out. push ( & self . on ) ;
348+ for clause in & self . clauses {
349+ if let Some ( predicate) = & clause. predicate {
350+ out. push ( predicate) ;
351+ }
352+ match & clause. action {
353+ MergeIntoAction :: Update ( assignments) => {
354+ out. extend ( assignments. iter ( ) . map ( |( _, value) | value) ) ;
355+ }
356+ MergeIntoAction :: Insert { values, .. } => {
357+ out. extend ( values. iter ( ) ) ;
358+ }
359+ MergeIntoAction :: Delete => { }
360+ }
361+ }
362+ out
363+ }
364+
365+ /// Rebuild this `MergeIntoOp` from a flat vector of new expressions, in
366+ /// the same order produced by [`Self::exprs`]. The clause kinds, action
367+ /// kinds, column lists, and presence/absence of each predicate are
368+ /// preserved from `self`.
369+ pub fn with_new_exprs ( & self , exprs : Vec < Expr > ) -> Result < Self > {
370+ let expected = self . expr_count ( ) ;
371+ if exprs. len ( ) != expected {
372+ return internal_err ! (
373+ "MergeIntoOp::with_new_exprs expected {expected} expressions, got {}" ,
374+ exprs. len( )
375+ ) ;
376+ }
377+ let mut iter = exprs. into_iter ( ) ;
378+ let on = iter. next ( ) . expect ( "non-empty by length check" ) ;
379+ let clauses = self
380+ . clauses
381+ . iter ( )
382+ . map ( |clause| {
383+ let predicate = clause
384+ . predicate
385+ . is_some ( )
386+ . then ( || iter. next ( ) . expect ( "non-empty by length check" ) ) ;
387+ let action = match & clause. action {
388+ MergeIntoAction :: Update ( assignments) => {
389+ let assignments = assignments
390+ . iter ( )
391+ . map ( |( name, _) | {
392+ (
393+ name. clone ( ) ,
394+ iter. next ( ) . expect ( "non-empty by length check" ) ,
395+ )
396+ } )
397+ . collect ( ) ;
398+ MergeIntoAction :: Update ( assignments)
399+ }
400+ MergeIntoAction :: Insert { columns, values } => {
401+ let values = values
402+ . iter ( )
403+ . map ( |_| iter. next ( ) . expect ( "non-empty by length check" ) )
404+ . collect ( ) ;
405+ MergeIntoAction :: Insert {
406+ columns : columns. clone ( ) ,
407+ values,
408+ }
409+ }
410+ MergeIntoAction :: Delete => MergeIntoAction :: Delete ,
411+ } ;
412+ MergeIntoClause {
413+ kind : clause. kind ,
414+ predicate,
415+ action,
416+ }
417+ } )
418+ . collect ( ) ;
419+ Ok ( Self {
420+ on,
421+ clauses,
422+ target_schema : Arc :: clone ( & self . target_schema ) ,
423+ } )
424+ }
308425}
309426
310427/// A single WHEN clause within a MERGE INTO statement.
@@ -400,6 +517,23 @@ fn make_count_schema() -> DFSchemaRef {
400517mod tests {
401518 use super :: * ;
402519 use crate :: { col, lit} ;
520+ use arrow:: datatypes:: DataType ;
521+ use datafusion_common:: DFSchema ;
522+
523+ fn dummy_target_schema ( ) -> DFSchemaRef {
524+ Arc :: new (
525+ DFSchema :: try_from_qualified_schema (
526+ "target" ,
527+ & Schema :: new ( vec ! [
528+ Field :: new( "id" , DataType :: Int64 , false ) ,
529+ Field :: new( "qty" , DataType :: Int64 , true ) ,
530+ Field :: new( "price" , DataType :: Float64 , true ) ,
531+ Field :: new( "active" , DataType :: Boolean , true ) ,
532+ ] ) ,
533+ )
534+ . unwrap ( ) ,
535+ )
536+ }
403537
404538 #[ test]
405539 fn write_op_merge_into_name_and_display ( ) {
@@ -413,6 +547,7 @@ mod tests {
413547 col( "source_qty" ) ,
414548 ) ] ) ,
415549 } ] ,
550+ target_schema : dummy_target_schema ( ) ,
416551 } ) ) ;
417552 assert_eq ! ( op. name( ) , "MergeInto" ) ;
418553 assert_eq ! ( format!( "{op}" ) , "MergeInto" ) ;
@@ -445,4 +580,55 @@ mod tests {
445580 MergeIntoClauseKind :: NotMatchedBySource
446581 ) ;
447582 }
583+
584+ #[ test]
585+ fn merge_into_op_exprs_round_trip ( ) {
586+ let op = MergeIntoOp {
587+ on : col ( "id" ) . eq ( col ( "source_id" ) ) ,
588+ clauses : vec ! [
589+ MergeIntoClause {
590+ kind: MergeIntoClauseKind :: Matched ,
591+ predicate: Some ( col( "qty" ) . gt( lit( 0_i64 ) ) ) ,
592+ action: MergeIntoAction :: Update ( vec![
593+ ( "qty" . to_string( ) , col( "source_qty" ) ) ,
594+ ( "price" . to_string( ) , col( "source_price" ) ) ,
595+ ] ) ,
596+ } ,
597+ MergeIntoClause {
598+ kind: MergeIntoClauseKind :: NotMatched ,
599+ predicate: None ,
600+ action: MergeIntoAction :: Insert {
601+ columns: vec![ "id" . to_string( ) , "qty" . to_string( ) ] ,
602+ values: vec![ col( "source_id" ) , col( "source_qty" ) ] ,
603+ } ,
604+ } ,
605+ MergeIntoClause {
606+ kind: MergeIntoClauseKind :: NotMatchedBySource ,
607+ predicate: Some ( col( "active" ) . eq( lit( true ) ) ) ,
608+ action: MergeIntoAction :: Delete ,
609+ } ,
610+ ] ,
611+ target_schema : dummy_target_schema ( ) ,
612+ } ;
613+ let exprs = op. exprs ( ) ;
614+ assert_eq ! ( exprs. len( ) , 7 ) ;
615+
616+ let owned: Vec < Expr > = exprs. into_iter ( ) . cloned ( ) . collect ( ) ;
617+ let rebuilt = op. with_new_exprs ( owned) . unwrap ( ) ;
618+ assert_eq ! ( op, rebuilt) ;
619+ }
620+
621+ #[ test]
622+ fn merge_into_op_with_new_exprs_length_mismatch ( ) {
623+ let op = MergeIntoOp {
624+ on : col ( "id" ) . eq ( col ( "source_id" ) ) ,
625+ clauses : vec ! [ ] ,
626+ target_schema : dummy_target_schema ( ) ,
627+ } ;
628+ let err = op. with_new_exprs ( vec ! [ ] ) . unwrap_err ( ) ;
629+ assert ! (
630+ err. to_string( ) . contains( "expected 1 expressions, got 0" ) ,
631+ "unexpected error: {err}"
632+ ) ;
633+ }
448634}
0 commit comments