Add related source code locations to errors

datafusion/common/src/diagnostic.rs

@@ -0,0 +1,85 @@
+use sqlparser::tokenizer::Span;
+
+#[derive(Debug, Clone)]
+pub struct Diagnostic {
+    pub entries: Vec<DiagnosticEntry>,
+}
+
+#[derive(Debug, Clone)]
+pub struct DiagnosticEntry {
+    pub span: Span,
+    pub message: String,
+    pub kind: DiagnosticEntryKind,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum DiagnosticEntryKind {
+    Error,
+    Warning,
+    Note,
+    Help,
+}
+
+impl Diagnostic {
+    pub fn new() -> Self {
+        Default::default()
+    }
+}
+
+impl Default for Diagnostic {
+    fn default() -> Self {
+        Diagnostic {
+            entries: Vec::new(),
+        }
+    }
+}
+
+impl FromIterator<DiagnosticEntry> for Diagnostic {
+    fn from_iter<T: IntoIterator<Item = DiagnosticEntry>>(iter: T) -> Self {
+        Diagnostic {
+            entries: iter.into_iter().collect(),
+        }
+    }
+}
+
+macro_rules! with_kind {
+    ($name:ident, $kind:expr) => {
+        pub fn $name(mut self, message: impl Into<String>, span: Span) -> Self {
+            let entry = DiagnosticEntry {
+                span,
+                message: message.into(),
+                kind: $kind,
+            };
+            self.entries.push(entry);
+            self
+        }
+    };
+}
+
+impl Diagnostic {
+    with_kind!(with_error, DiagnosticEntryKind::Error);
+    with_kind!(with_warning, DiagnosticEntryKind::Warning);
+    with_kind!(with_note, DiagnosticEntryKind::Note);
+    with_kind!(with_help, DiagnosticEntryKind::Help);
+}
+
+impl DiagnosticEntry {
+    pub fn new(
+        message: impl Into<String>,
+        kind: DiagnosticEntryKind,
+        span: Span,
+    ) -> Self {
+        DiagnosticEntry {
+            span,
+            message: message.into(),
+            kind,
+        }
+    }
+
+    pub fn new_without_span(
+        message: impl Into<String>,
+        kind: DiagnosticEntryKind,
+    ) -> Self {
+        Self::new(message, kind, Span::empty())
+    }
+}

datafusion/common/src/error.rs

@@ -27,7 +27,7 @@ use std::result;
 use std::sync::Arc;
 
 use crate::utils::quote_identifier;
-use crate::{Column, DFSchema, TableReference};
+use crate::{Column, DFSchema, Diagnostic, TableReference};
 #[cfg(feature = "avro")]
 use apache_avro::Error as AvroError;
 use arrow::error::ArrowError;
@@ -131,6 +131,11 @@ pub enum DataFusionError {
     /// Errors from either mapping LogicalPlans to/from Substrait plans
     /// or serializing/deserializing protobytes to Substrait plans
     Substrait(String),
+    /// Error wrapped together with additional contextual information intended
+    /// for end users, to help them understand what went wrong by providing
+    /// human-readable messages, and locations in the source query that relate
+    /// to the error in some way.
+    Diagnostic(Diagnostic, Box<DataFusionError>),
 }
 
 #[macro_export]
@@ -328,6 +333,7 @@ impl Error for DataFusionError {
             DataFusionError::External(e) => Some(e.as_ref()),
             DataFusionError::Context(_, e) => Some(e.as_ref()),
             DataFusionError::Substrait(_) => None,
+            DataFusionError::Diagnostic(_, e) => Some(e.as_ref()),
         }
     }
 }
@@ -441,6 +447,7 @@ impl DataFusionError {
             DataFusionError::External(_) => "External error: ",
             DataFusionError::Context(_, _) => "",
             DataFusionError::Substrait(_) => "Substrait error: ",
+            DataFusionError::Diagnostic(_, _) => "",
         }
     }
 
@@ -481,8 +488,56 @@ impl DataFusionError {
                 Cow::Owned(format!("{desc}\ncaused by\n{}", *err))
             }
             DataFusionError::Substrait(ref desc) => Cow::Owned(desc.to_string()),
+            DataFusionError::Diagnostic(_, ref err) => Cow::Owned(err.to_string()),
         }
     }
+
+    /// Wraps the error with contextual information intended for end users
+    pub fn with_diagnostic(self, diagnostic: Diagnostic) -> Self {
+        Self::Diagnostic(diagnostic, Box::new(self))
+    }
+
+    /// Wraps the error with contextual information intended for end users.
+    /// Takes a function that inspects the error and returns the diagnostic to
+    /// wrap it with.
+    pub fn with_diagnostic_fn<F: FnOnce(&DataFusionError) -> Diagnostic>(
+        self,
+        f: F,
+    ) -> Self {
+        let diagnostic = f(&self);
+        self.with_diagnostic(diagnostic)
+    }
+
+    pub fn get_diagnostics(&self) -> impl Iterator<Item = &Diagnostic> + '_ {
+        struct DiagnosticsIterator<'a> {
+            head: &'a DataFusionError,
+        }
+
+        impl<'a> Iterator for DiagnosticsIterator<'a> {
+            type Item = &'a Diagnostic;
+
+            fn next(&mut self) -> Option<Self::Item> {
+                loop {
+                    if let DataFusionError::Diagnostic(diagnostics, source) = self.head {
+                        self.head = source.as_ref();
+                        return Some(diagnostics);
+                    }
+
+                    if let Some(source) = self
+                        .head
+                        .source()
+                        .and_then(|source| source.downcast_ref::<DataFusionError>())
+                    {
+                        self.head = source;
+                    } else {
+                        return None;
+                    }
+                }
+            }
+        }
+
+        DiagnosticsIterator { head: self }
+    }
 }
 
 /// Unwrap an `Option` if possible. Otherwise return an `DataFusionError::Internal`.

datafusion/common/src/lib.rs

@@ -47,6 +47,7 @@ pub mod test_util;
 pub mod tree_node;
 pub mod types;
 pub mod utils;
+pub mod diagnostic;
 
 /// Reexport arrow crate
 pub use arrow;
@@ -76,6 +77,7 @@ pub use stats::{ColumnStatistics, Statistics};
 pub use table_reference::{ResolvedTableReference, TableReference};
 pub use unnest::{RecursionUnnestOption, UnnestOptions};
 pub use utils::project_schema;
+pub use diagnostic::{Diagnostic, DiagnosticEntry, DiagnosticEntryKind};
 
 // These are hidden from docs purely to avoid polluting the public view of what this crate exports.
 // These are just re-exports of macros by the same name, which gets around the 'cannot refer to

datafusion/sql/src/expr/mod.rs

@@ -15,11 +15,21 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::ops::ControlFlow;
+
 use arrow_schema::DataType;
 use arrow_schema::TimeUnit;
+use datafusion_common::Column;
+use datafusion_common::DataFusionError;
+use datafusion_common::Diagnostic;
+use datafusion_common::SchemaError;
 use datafusion_expr::planner::{
     PlannerResult, RawBinaryExpr, RawDictionaryExpr, RawFieldAccessExpr,
 };
+use datafusion_expr::utils::find_column_exprs;
+use sqlparser::ast::Spanned;
+use sqlparser::ast::Visit;
+use sqlparser::ast::Visitor;
 use sqlparser::ast::{
     BinaryOperator, CastFormat, CastKind, DataType as SQLDataType, DictionaryField,
     Expr as SQLExpr, ExprWithAlias as SQLExprWithAlias, MapEntry, StructField, Subscript,
@@ -36,7 +46,9 @@ use datafusion_expr::{
     lit, Between, BinaryExpr, Cast, Expr, ExprSchemable, GetFieldAccess, Like, Literal,
     Operator, TryCast,
 };
+use sqlparser::tokenizer::Span;
 
+use crate::planner::IdentNormalizer;
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
 
 mod binary_op;
@@ -165,13 +177,126 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         schema: &DFSchema,
         planner_context: &mut PlannerContext,
     ) -> Result<Expr> {
-        let mut expr = self.sql_expr_to_logical_expr(sql, schema, planner_context)?;
+        // The location of the original SQL expression in the source code
+        let mut expr =
+            self.sql_expr_to_logical_expr(sql.clone(), schema, planner_context)?;
         expr = self.rewrite_partial_qualifier(expr, schema);
-        self.validate_schema_satisfies_exprs(schema, std::slice::from_ref(&expr))?;
+        let validation_result =
+            self.validate_schema_satisfies_exprs(schema, std::slice::from_ref(&expr));
+
+        // Must do it here because `validate_schema_satisfies_exprs` doesn't
+        // have access to the original SQL expression from the parser
+        validation_result.map_err(|err| {
+            if let DataFusionError::SchemaError(
+                SchemaError::FieldNotFound {
+                    field,
+                    valid_fields: _,
+                },
+                _,
+            ) = &err
+            {
+                let diagnostic = self.get_field_not_found_diagnostic(
+                    &sql,
+                    field,
+                    schema,
+                    planner_context,
+                );
+                err.with_diagnostic(diagnostic)
+            } else {
+                err
+            }
+        })?;
+
         let (expr, _) = expr.infer_placeholder_types(schema)?;
         Ok(expr)
     }
 
+    /// Given an unresolved field in an expression, returns a [`Diagnostic`]
+    fn get_field_not_found_diagnostic(
+        &self,
+        expr: &SQLExpr,
+        unresolved_field: &Column,
+        schema: &DFSchema,
+        planner_context: &mut PlannerContext,
+    ) -> Diagnostic {
+        // Given a SQL expression like SELECT color, megabytes FROM fruit, where
+        // we assume that the 'megabytes' column doesn't exist in table 'fruit',
+        // we find that the logical expression Expr::Column(Column('megabytes'))
+        // is unresolved. Though, because we don't store `sqlparser::Span` in
+        // `datafusion::Expr`, we have no simple way to find where 'megabytes'
+        // appears in the query.
+        //
+        // Instead, we have to walk down the tree of subexpressions
+        // `sqlparser::Expr` rooted in the parameter `expr: sqlparser::Expr` to
+        // find a node that matches the unresolved one.
+        struct UnresolvedFieldFinder<'a, S: ContextProvider> {
+            unresolved_field: &'a Column,
+            sql_to_rel: &'a SqlToRel<'a, S>,
+            schema: &'a DFSchema,
+            planner_context: &'a mut PlannerContext,
+        }
+        impl<S: ContextProvider> Visitor for UnresolvedFieldFinder<'_, S> {
+            type Break = Span;
+
+            fn post_visit_expr(
+                &mut self,
+                sql_expr: &SQLExpr,
+            ) -> ControlFlow<Self::Break> {
+                let (logical_expr, span) = match sql_expr {
+                    SQLExpr::Identifier(ident) => (
+                        self.sql_to_rel
+                            .sql_identifier_to_expr(
+                                ident.clone(),
+                                self.schema,
+                                self.planner_context,
+                            )
+                            .ok(),
+                        sql_expr.span(),
+                    ),
+                    SQLExpr::CompoundIdentifier(idents) => (
+                        self.sql_to_rel
+                            .sql_compound_identifier_to_expr(
+                                idents.clone(),
+                                self.schema,
+                                self.planner_context,
+                            )
+                            .ok(),
+                        sql_expr.span(),
+                    ),
+                    _ => (None, Span::empty()),
+                };
+                match logical_expr {
+                    Some(Expr::Column(col)) if &col == self.unresolved_field => {
+                        ControlFlow::Break(span)
+                    }
+                    _ => ControlFlow::Continue(()),
+                }
+            }
+        }
+        let mut visitor = UnresolvedFieldFinder {
+            unresolved_field: &unresolved_field,
+            sql_to_rel: self,
+            schema,
+            planner_context,
+        };
+        let span = match expr.visit(&mut visitor) {
+            ControlFlow::Break(span) => Some(span),
+            ControlFlow::Continue(_) => None,
+        };
+
+        if let Some(relation) = &unresolved_field.relation {
+            Diagnostic::new().with_error(
+                format!("column '{}' not found in '{}'", &unresolved_field.name, relation.to_string()),
+                span.unwrap_or(Span::empty()),
+            )
+        } else {
+            Diagnostic::new().with_error(
+                format!("column '{}' not found", &unresolved_field.name),
+                span.unwrap_or(Span::empty()),
+            )
+        }
+    }
+
     /// Rewrite aliases which are not-complete (e.g. ones that only include only table qualifier in a schema.table qualified relation)
     fn rewrite_partial_qualifier(&self, expr: Expr, schema: &DFSchema) -> Expr {
         match expr {

datafusion/sql/src/parser.rs

@@ -284,10 +284,10 @@ impl<'a> DFParser<'a> {
         dialect: &'a dyn Dialect,
     ) -> Result<Self, ParserError> {
         let mut tokenizer = Tokenizer::new(dialect, sql);
-        let tokens = tokenizer.tokenize()?;
+        let tokens = tokenizer.tokenize_with_location()?;
 
         Ok(DFParser {
-            parser: Parser::new(dialect).with_tokens(tokens),
+            parser: Parser::new(dialect).with_tokens_with_locations(tokens),
         })
     }