fix aggr/window plan

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

Author: waynexia

datafusion/core/tests/set_comparison.rs

@@ -55,6 +55,24 @@ async fn set_comparison_any() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn set_comparison_any_aggregate_subquery() -> Result<()> {
+    let ctx = SessionContext::new();
+
+    ctx.register_batch("t", build_table(&[1, 7])?)?;
+    ctx.register_batch("s", build_table(&[1, 2, 3])?)?;
+
+    let df = ctx
+        .sql(
+            "select v from t where v > any(select sum(v) from s group by v % 2) order by v",
+        )
+        .await?;
+    let results = df.collect().await?;
+
+    assert_batches_eq!(&["+---+", "| v |", "+---+", "| 7 |", "+---+",], &results);
+    Ok(())
+}
+
 #[tokio::test]
 async fn set_comparison_all_empty() -> Result<()> {
     let ctx = SessionContext::new();

datafusion/optimizer/src/rewrite_set_comparison.rs

@@ -20,9 +20,8 @@
 //! that capture SQL three-valued logic.
 
 use crate::{OptimizerConfig, OptimizerRule};
-use datafusion_common::ExprSchema;
 use datafusion_common::tree_node::{Transformed, TreeNode};
-use datafusion_common::{DFSchema, Result, ScalarValue, plan_datafusion_err};
+use datafusion_common::{Column, DFSchema, ExprSchema, Result, ScalarValue, plan_err};
 use datafusion_expr::expr::{self, Exists, SetComparison, SetQuantifier};
 use datafusion_expr::logical_plan::Subquery;
 use datafusion_expr::logical_plan::builder::LogicalPlanBuilder;
@@ -91,10 +90,12 @@ fn build_set_comparison_subquery(
     } = set_comparison;
 
     let left_expr = to_outer_reference(*expr, outer_schema)?;
-    let right_expr = subquery
-        .subquery
-        .head_output_expr()?
-        .ok_or_else(|| plan_datafusion_err!("single expression required."))?;
+    let subquery_schema = subquery.subquery.schema();
+    if subquery_schema.fields().is_empty() {
+        return plan_err!("single expression required.");
+    }
+    // avoid `head_output_expr` for aggr/window plan, it will gives group-by expr if exists
+    let right_expr = Expr::Column(Column::from(subquery_schema.qualified_field(0)));
 
     let comparison = Expr::BinaryExpr(expr::BinaryExpr::new(
         Box::new(left_expr),