mingmwang commented on code in PR #6084:
URL: https://github.com/apache/arrow-datafusion/pull/6084#discussion_r1177516037
##########
datafusion/core/tests/sql/subqueries.rs:
##########
@@ -203,7 +203,202 @@ async fn subquery_not_allowed() -> Result<()> {
let err = dataframe.into_optimized_plan().err().unwrap();
assert_eq!(
- r#"Context("check_analyzed_plan", Plan("In/Exist subquery can not be
used in Sort plan nodes"))"#,
+ r#"Context("check_analyzed_plan", Plan("In/Exist subquery can only be
used in Projection, Filter, Window functions, Aggregate and Join plan
nodes"))"#,
+ &format!("{err:?}")
+ );
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn non_aggregated_correlated_scalar_subquery() -> Result<()> {
+ let ctx = create_join_context("t1_id", "t2_id", true)?;
+
+ let sql = "SELECT t1_id, (SELECT t2_int FROM t2 WHERE t2.t2_int =
t1.t1_int) as t2_int from t1";
+ let msg = format!("Creating logical plan for '{sql}'");
+ let dataframe = ctx.sql(sql).await.expect(&msg);
+ let err = dataframe.into_optimized_plan().err().unwrap();
+
+ assert_eq!(
+ r#"Context("check_analyzed_plan", Plan("Correlated scalar subquery
must be aggregated to return at most one row"))"#,
+ &format!("{err:?}")
+ );
+
+ let sql = "SELECT t1_id, (SELECT t2_int FROM t2 WHERE t2.t2_int = t1_int
group by t2_int) as t2_int from t1";
+ let msg = format!("Creating logical plan for '{sql}'");
+ let dataframe = ctx.sql(sql).await.expect(&msg);
+ let err = dataframe.into_optimized_plan().err().unwrap();
+
+ assert_eq!(
+ r#"Context("check_analyzed_plan", Plan("Correlated scalar subquery
must be aggregated to return at most one row"))"#,
+ &format!("{err:?}")
+ );
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn non_aggregated_correlated_scalar_subquery_with_limit() -> Result<()> {
+ let ctx = create_join_context("t1_id", "t2_id", true)?;
+
+ let sql = "SELECT t1_id, (SELECT t2_int FROM t2 WHERE t2.t2_int =
t1.t1_int limit 2) as t2_int from t1";
+ let msg = format!("Creating logical plan for '{sql}'");
+ let dataframe = ctx.sql(sql).await.expect(&msg);
+ let err = dataframe.into_optimized_plan().err().unwrap();
+
+ assert_eq!(
+ r#"Context("check_analyzed_plan", Plan("Correlated scalar subquery
must be aggregated to return at most one row"))"#,
+ &format!("{err:?}")
+ );
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn non_aggregated_correlated_scalar_subquery_with_single_row() ->
Result<()> {
+ let ctx = create_join_context("t1_id", "t2_id", true)?;
+
+ let sql = "SELECT t1_id, (SELECT t2_int FROM t2 WHERE t2.t2_int =
t1.t1_int limit 1) as t2_int from t1";
+ let msg = format!("Creating logical plan for '{sql}'");
+ let dataframe = ctx.sql(sql).await.expect(&msg);
+ let plan = dataframe.into_optimized_plan()?;
+
+ let expected = vec![
+ "Projection: t1.t1_id, (<subquery>) AS t2_int [t1_id:UInt32;N,
t2_int:UInt32;N]",
+ " Subquery: [t2_int:UInt32;N]",
+ " Limit: skip=0, fetch=1 [t2_int:UInt32;N]",
+ " Projection: t2.t2_int [t2_int:UInt32;N]",
+ " Filter: t2.t2_int = outer_ref(t1.t1_int) [t2_id:UInt32;N,
t2_name:Utf8;N, t2_int:UInt32;N]",
+ " TableScan: t2 [t2_id:UInt32;N, t2_name:Utf8;N,
t2_int:UInt32;N]",
+ " TableScan: t1 projection=[t1_id] [t1_id:UInt32;N]",
+ ];
+ let formatted = plan.display_indent_schema().to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+ assert_eq!(
+ expected, actual,
+ "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ let sql = "SELECT t1_id from t1 where t1_int = (SELECT t2_int FROM t2
WHERE t2.t2_int = t1.t1_int limit 1)";
+ let msg = format!("Creating logical plan for '{sql}'");
+ let dataframe = ctx.sql(sql).await.expect(&msg);
+ let plan = dataframe.into_optimized_plan()?;
+
+ let expected = vec![
+ "Projection: t1.t1_id [t1_id:UInt32;N]",
+ " Filter: t1.t1_int = (<subquery>) [t1_id:UInt32;N, t1_int:UInt32;N]",
+ " Subquery: [t2_int:UInt32;N]",
+ " Limit: skip=0, fetch=1 [t2_int:UInt32;N]",
+ " Projection: t2.t2_int [t2_int:UInt32;N]",
+ " Filter: t2.t2_int = outer_ref(t1.t1_int) [t2_id:UInt32;N,
t2_name:Utf8;N, t2_int:UInt32;N]",
+ " TableScan: t2 [t2_id:UInt32;N, t2_name:Utf8;N,
t2_int:UInt32;N]",
+ " TableScan: t1 projection=[t1_id, t1_int] [t1_id:UInt32;N,
t1_int:UInt32;N]",
+ ];
+ let formatted = plan.display_indent_schema().to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+ assert_eq!(
+ expected, actual,
+ "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ let sql = "SELECT t1_id, (SELECT a FROM (select 1 as a) WHERE a =
t1.t1_int) as t2_int from t1";
+ let msg = format!("Creating logical plan for '{sql}'");
+ let dataframe = ctx.sql(sql).await.expect(&msg);
+ let plan = dataframe.into_optimized_plan()?;
+
+ let expected = vec![
+ "Projection: t1.t1_id, (<subquery>) AS t2_int [t1_id:UInt32;N,
t2_int:Int64]",
+ " Subquery: [a:Int64]",
+ " Projection: a [a:Int64]",
+ " Filter: a = CAST(outer_ref(t1.t1_int) AS Int64) [a:Int64]",
+ " Projection: Int64(1) AS a [a:Int64]",
+ " EmptyRelation []",
+ " TableScan: t1 projection=[t1_id] [t1_id:UInt32;N]",
+ ];
+ let formatted = plan.display_indent_schema().to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+ assert_eq!(
+ expected, actual,
+ "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn non_equal_correlated_scalar_subquery() -> Result<()> {
+ let ctx = create_join_context("t1_id", "t2_id", true)?;
+
+ let sql = "SELECT t1_id, (SELECT sum(t2_int) FROM t2 WHERE t2.t2_id <
t1.t1_id) as t2_sum from t1";
+ let msg = format!("Creating logical plan for '{sql}'");
+ let dataframe = ctx.sql(sql).await.expect(&msg);
+ let err = dataframe.into_optimized_plan().err().unwrap();
+
+ assert_eq!(
+ r#"Context("check_analyzed_plan", Plan("Correlated column is not
allowed in predicate: t2.t2_id < outer_ref(t1.t1_id)"))"#,
+ &format!("{err:?}")
+ );
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn aggregated_correlated_scalar_subquery() -> Result<()> {
+ let ctx = create_join_context("t1_id", "t2_id", true)?;
+
+ let sql = "SELECT t1_id, (SELECT sum(t2_int) FROM t2 WHERE t2.t2_id =
t1.t1_id) as t2_sum from t1";
+ let msg = format!("Creating logical plan for '{sql}'");
+ let dataframe = ctx.sql(sql).await.expect(&msg);
+ let plan = dataframe.into_optimized_plan()?;
+
+ let expected = vec![
+ "Projection: t1.t1_id, (<subquery>) AS t2_sum [t1_id:UInt32;N,
t2_sum:UInt64;N]",
+ " Subquery: [SUM(t2.t2_int):UInt64;N]",
+ " Projection: SUM(t2.t2_int) [SUM(t2.t2_int):UInt64;N]",
+ " Aggregate: groupBy=[[]], aggr=[[SUM(t2.t2_int)]]
[SUM(t2.t2_int):UInt64;N]",
+ " Filter: t2.t2_id = outer_ref(t1.t1_id) [t2_id:UInt32;N,
t2_name:Utf8;N, t2_int:UInt32;N]",
+ " TableScan: t2 [t2_id:UInt32;N, t2_name:Utf8;N,
t2_int:UInt32;N]",
+ " TableScan: t1 projection=[t1_id] [t1_id:UInt32;N]",
+ ];
+ let formatted = plan.display_indent_schema().to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+ assert_eq!(
+ expected, actual,
+ "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn aggregated_correlated_scalar_subquery_with_extra_group_by_columns()
-> Result<()>
+{
+ let ctx = create_join_context("t1_id", "t2_id", true)?;
+
+ let sql = "SELECT t1_id, (SELECT sum(t2_int) FROM t2 WHERE t2.t2_id =
t1.t1_id group by t2_name) as t2_sum from t1";
+ let msg = format!("Creating logical plan for '{sql}'");
+ let dataframe = ctx.sql(sql).await.expect(&msg);
+ let err = dataframe.into_optimized_plan().err().unwrap();
+
+ assert_eq!(
+ r#"Context("check_analyzed_plan", Plan("A GROUP BY clause in a scalar
correlated subquery cannot contain non-correlated columns"))"#,
+ &format!("{err:?}")
+ );
+
+ Ok(())
+}
+
Review Comment:
Done
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
