@@ -699,15 +699,13 @@ impl LogicalPlan {
699
699
} ) )
700
700
}
701
701
LogicalPlan :: Union ( Union { inputs, schema } ) => {
702
- let input_schema = inputs[ 0 ] . schema ( ) ;
703
- // If inputs are not pruned do not change schema
704
- // TODO this seems wrong (shouldn't we always use the schema of the input?)
705
- let schema = if schema. fields ( ) . len ( ) == input_schema. fields ( ) . len ( ) {
706
- Arc :: clone ( & schema)
702
+ let first_input_schema = inputs[ 0 ] . schema ( ) ;
703
+ if schema. fields ( ) . len ( ) == first_input_schema. fields ( ) . len ( ) {
704
+ // If inputs are not pruned do not change schema
705
+ Ok ( LogicalPlan :: Union ( Union { inputs, schema } ) )
707
706
} else {
708
- Arc :: clone ( input_schema)
709
- } ;
710
- Ok ( LogicalPlan :: Union ( Union { inputs, schema } ) )
707
+ Ok ( LogicalPlan :: Union ( Union :: try_new ( inputs) ?) )
708
+ }
711
709
}
712
710
LogicalPlan :: Distinct ( distinct) => {
713
711
let distinct = match distinct {
@@ -2645,6 +2643,106 @@ pub struct Union {
2645
2643
pub schema : DFSchemaRef ,
2646
2644
}
2647
2645
2646
+ impl Union {
2647
+ /// Constructs new Union instance deriving schema from inputs.
2648
+ fn try_new ( inputs : Vec < Arc < LogicalPlan > > ) -> Result < Self > {
2649
+ let schema = Self :: derive_schema_from_inputs ( & inputs, false ) ?;
2650
+ Ok ( Union { inputs, schema } )
2651
+ }
2652
+
2653
+ /// Constructs new Union instance deriving schema from inputs.
2654
+ /// Inputs do not have to have matching types and produced schema will
2655
+ /// take type from the first input.
2656
+ pub fn try_new_with_loose_types ( inputs : Vec < Arc < LogicalPlan > > ) -> Result < Self > {
2657
+ let schema = Self :: derive_schema_from_inputs ( & inputs, true ) ?;
2658
+ Ok ( Union { inputs, schema } )
2659
+ }
2660
+
2661
+ /// Constructs new Union instance deriving schema from inputs.
2662
+ ///
2663
+ /// `loose_types` if true, inputs do not have to have matching types and produced schema will
2664
+ /// take type from the first input. TODO this is not necessarily reasonable behavior.
2665
+ fn derive_schema_from_inputs (
2666
+ inputs : & [ Arc < LogicalPlan > ] ,
2667
+ loose_types : bool ,
2668
+ ) -> Result < DFSchemaRef > {
2669
+ if inputs. len ( ) < 2 {
2670
+ return plan_err ! ( "UNION requires at least two inputs" ) ;
2671
+ }
2672
+ let first_schema = inputs[ 0 ] . schema ( ) ;
2673
+ let fields_count = first_schema. fields ( ) . len ( ) ;
2674
+ for input in inputs {
2675
+ if fields_count != input. schema ( ) . fields ( ) . len ( ) {
2676
+ return plan_err ! (
2677
+ "UNION queries have different number of columns: \
2678
+ left has {} columns whereas right has {} columns",
2679
+ fields_count,
2680
+ input. schema( ) . fields( ) . len( )
2681
+ ) ;
2682
+ }
2683
+ }
2684
+
2685
+ let union_fields = ( 0 ..fields_count)
2686
+ . map ( |i| {
2687
+ let fields = inputs
2688
+ . iter ( )
2689
+ . map ( |input| input. schema ( ) . field ( i) )
2690
+ . collect :: < Vec < _ > > ( ) ;
2691
+ let first_field = fields[ 0 ] ;
2692
+ let name = first_field. name ( ) ;
2693
+ let data_type = if loose_types {
2694
+ // TODO apply type coercion here, or document why it's better to defer
2695
+ // temporarily use the data type from the left input and later rely on the analyzer to
2696
+ // coerce the two schemas into a common one.
2697
+ first_field. data_type ( )
2698
+ } else {
2699
+ fields. iter ( ) . skip ( 1 ) . try_fold (
2700
+ first_field. data_type ( ) ,
2701
+ |acc, field| {
2702
+ if acc != field. data_type ( ) {
2703
+ return plan_err ! (
2704
+ "UNION field {i} have different type in inputs: \
2705
+ left has {} whereas right has {}",
2706
+ first_field. data_type( ) ,
2707
+ field. data_type( )
2708
+ ) ;
2709
+ }
2710
+ Ok ( acc)
2711
+ } ,
2712
+ ) ?
2713
+ } ;
2714
+ let nullable = fields. iter ( ) . any ( |field| field. is_nullable ( ) ) ;
2715
+ let mut field = Field :: new ( name, data_type. clone ( ) , nullable) ;
2716
+ let field_metadata =
2717
+ intersect_maps ( fields. iter ( ) . map ( |field| field. metadata ( ) ) ) ;
2718
+ field. set_metadata ( field_metadata) ;
2719
+ // TODO reusing table reference from the first schema is probably wrong
2720
+ let table_reference = first_schema. qualified_field ( i) . 0 . cloned ( ) ;
2721
+ Ok ( ( table_reference, Arc :: new ( field) ) )
2722
+ } )
2723
+ . collect :: < Result < _ > > ( ) ?;
2724
+ let union_schema_metadata =
2725
+ intersect_maps ( inputs. iter ( ) . map ( |input| input. schema ( ) . metadata ( ) ) ) ;
2726
+
2727
+ // Functional Dependencies doesn't preserve after UNION operation
2728
+ let schema = DFSchema :: new_with_metadata ( union_fields, union_schema_metadata) ?;
2729
+ let schema = Arc :: new ( schema) ;
2730
+
2731
+ Ok ( schema)
2732
+ }
2733
+ }
2734
+
2735
+ fn intersect_maps < ' a > (
2736
+ inputs : impl IntoIterator < Item = & ' a HashMap < String , String > > ,
2737
+ ) -> HashMap < String , String > {
2738
+ let mut inputs = inputs. into_iter ( ) ;
2739
+ let mut merged: HashMap < String , String > = inputs. next ( ) . cloned ( ) . unwrap_or_default ( ) ;
2740
+ for input in inputs {
2741
+ merged. retain ( |k, v| input. get ( k) == Some ( v) ) ;
2742
+ }
2743
+ merged
2744
+ }
2745
+
2648
2746
// Manual implementation needed because of `schema` field. Comparison excludes this field.
2649
2747
impl PartialOrd for Union {
2650
2748
fn partial_cmp ( & self , other : & Self ) -> Option < Ordering > {
0 commit comments