changsun20 commented on code in PR #15143:
URL: https://github.com/apache/datafusion/pull/15143#discussion_r1990538405


##########
datafusion/sql/tests/cases/diagnostic.rs:
##########
@@ -286,3 +286,121 @@ fn test_invalid_function() -> Result<()> {
     assert_eq!(diag.span, Some(spans["whole"]));
     Ok(())
 }
+
+#[test]
+fn test_scalar_subquery_multiple_columns() -> Result<(), Box<dyn 
std::error::Error>> {
+    let query = "SELECT (SELECT 1 AS /*x*/x/*x*/, 2 AS /*y*/y/*y*/) AS col";
+    let spans = get_spans(query);
+    let diag = do_query(query);
+
+    assert_eq!(
+        diag.message,
+        "Scalar subquery should only return one column"
+    );
+
+    let column_count_notes: Vec<_> = diag
+        .notes
+        .iter()
+        .filter(|note| note.message.contains("Found 2 columns"))
+        .collect();
+    assert!(
+        !column_count_notes.is_empty(),
+        "Expected note about column count"
+    );
+
+    let extra_column_notes: Vec<_> = diag
+        .notes
+        .iter()
+        .filter(|note| note.message.contains("Extra column"))
+        .collect();
+    assert_eq!(
+        extra_column_notes.len(),
+        1,
+        "Expected one note about extra column"
+    );
+
+    if let Some(first_note) = diag.notes.first() {
+        assert_eq!(
+            first_note.span,
+            Some(spans["x"]),
+            "Primary diagnostic span should match the first column's span"
+        );
+    }
+
+    if let Some(extra_column_span) = extra_column_notes
+        .first()
+        .and_then(|note| note.span.as_ref())
+    {
+        assert_eq!(
+            *extra_column_span, spans["y"],
+            "Extra column note span should match the second column's span"
+        );
+    }

Review Comment:
   Changed. Thanks for the suggestion, this is indeed a more elegant way to 
test.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org

Reply via email to