feat: UNION wrong number of fields diagnostics

Author: eliaperantoni

datafusion/expr/src/logical_plan/builder.rs

@@ -52,7 +52,11 @@ use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef};
 use datafusion_common::display::ToStringifiedPlan;
 use datafusion_common::file_options::file_type::FileType;
 use datafusion_common::{
-    exec_err, get_target_functional_dependencies, internal_err, not_impl_err, plan_datafusion_err, plan_err, Column, DFSchema, DFSchemaRef, DataFusionError, FieldsSpans, FunctionalDependencies, Result, ScalarValue, TableReference, ToDFSchema, UnnestOptions
+    exec_err, get_target_functional_dependencies, internal_err, not_impl_err,
+    plan_datafusion_err, plan_err, Column, DFSchema, DFSchemaRef, DataFusionError,
+    Diagnostic, DiagnosticEntry, DiagnosticEntryKind, FieldsSpans,
+    FunctionalDependencies, Result, ScalarValue, TableReference, ToDFSchema,
+    UnnestOptions,
 };
 use datafusion_expr_common::type_coercion::binary::type_union_resolution;
 
@@ -1436,9 +1440,12 @@ pub fn build_join_schema(
         join_type,
         left.fields().len(),
     );
-    let left_fields_spans: FieldsSpans = left.iter().map(|(_, _, spans)| spans).cloned().collect();
-    let right_fields_spans: FieldsSpans = right.iter().map(|(_, _, spans)| spans).cloned().collect();
-    let fields_spans = left_fields_spans.join(&right_fields_spans, join_type, left.fields().len());
+    let left_fields_spans: FieldsSpans =
+        left.iter().map(|(_, _, spans)| spans).cloned().collect();
+    let right_fields_spans: FieldsSpans =
+        right.iter().map(|(_, _, spans)| spans).cloned().collect();
+    let fields_spans =
+        left_fields_spans.join(&right_fields_spans, join_type, left.fields().len());
     let metadata = left
         .metadata()
         .clone()
@@ -1526,7 +1533,32 @@ pub fn union(left_plan: LogicalPlan, right_plan: LogicalPlan) -> Result<LogicalP
             left has {} columns whereas right has {} columns",
             left_plan.schema().fields().len(),
             right_plan.schema().fields().len()
-        );
+        )
+        .map_err(|err| {
+            err.with_diagnostic(|_| {
+                Diagnostic::new([
+                    DiagnosticEntry::new(
+                        "UNION queries have different number of columns",
+                        DiagnosticEntryKind::Error,
+                        Span::empty(),
+                    ),
+                    DiagnosticEntry::new(
+                        format!("This side has {} columns", left_plan.schema().fields().len()),
+                        DiagnosticEntryKind::Note,
+                        Span::union_iter(
+                            left_plan.schema().fields_spans().iter().flatten().copied(),
+                        ),
+                    ),
+                    DiagnosticEntry::new(
+                        format!("This side has {} columns", right_plan.schema().fields().len()),
+                        DiagnosticEntryKind::Note,
+                        Span::union_iter(
+                            right_plan.schema().fields_spans().iter().flatten().copied(),
+                        ),
+                    ),
+                ])
+            })
+        });
     }
 
     // Temporarily use the schema from the left input and later rely on the analyzer to