Merge pull request #5662 from leiysky/having-alias

improve(planner): Support use alias group item in `HAVING`

query/src/sql/planner/binder/aggregate.rs

@@ -241,27 +241,37 @@ impl<'a> Binder {
     ///     `SELECT a as b, COUNT(a) FROM t GROUP BY b`.
     ///   - Scalar expressions that can be evaluated in current scope(doesn't contain aliases), e.g.
     ///     column `a` and expression `a+1` in `SELECT a as b, COUNT(a) FROM t GROUP BY a, a+1`.
-    pub(super) async fn bind_aggregate(
+    pub async fn analyze_group_items(
         &mut self,
         bind_context: &mut BindContext,
         select_list: &SelectList<'a>,
         group_by: &[Expr<'a>],
-        child: SExpr,
-    ) -> Result<SExpr> {
+    ) -> Result<()> {
         let mut available_aliases = vec![];
 
         // Extract available aliases from `SELECT` clause,
         for item in select_list.items.iter() {
             if let SelectTarget::AliasedExpr { alias: Some(_), .. } = item.select_target {
-                let column =
-                    self.create_column_binding(None, item.alias.clone(), item.scalar.data_type());
+                let column = if let Scalar::BoundColumnRef(column_ref) = &item.scalar {
+                    let mut column = column_ref.column.clone();
+                    column.column_name = item.alias.clone();
+                    column
+                } else {
+                    self.create_column_binding(None, item.alias.clone(), item.scalar.data_type())
+                };
                 available_aliases.push((column, item.scalar.clone()));
             }
         }
 
         self.resolve_group_items(bind_context, select_list, group_by, &available_aliases)
-            .await?;
+            .await
+    }
 
+    pub(super) async fn bind_aggregate(
+        &mut self,
+        bind_context: &mut BindContext,
+        child: SExpr,
+    ) -> Result<SExpr> {
         // Enter in_grouping state
         bind_context.in_grouping = true;
 
@@ -354,7 +364,7 @@ impl<'a> Binder {
             let (scalar_expr, data_type) = scalar_binder
                 .bind(expr)
                 .await
-                .or_else(|e| Self::resolve_alias_item(expr, available_aliases, e))?;
+                .or_else(|e| Self::resolve_alias_item(bind_context, expr, available_aliases, e))?;
 
             if bind_context
                 .aggregate_info
@@ -420,18 +430,25 @@ impl<'a> Binder {
     }
 
     fn resolve_alias_item(
+        bind_context: &mut BindContext,
         expr: &Expr<'a>,
         available_aliases: &[(ColumnBinding, Scalar)],
         original_error: ErrorCode,
     ) -> Result<(Scalar, DataTypeImpl)> {
-        let mut result: Vec<Scalar> = vec![];
+        let mut result: Vec<usize> = vec![];
         // If cannot resolve group item, then try to find an available alias
-        for (column_binding, scalar) in available_aliases.iter() {
+        for (i, (column_binding, _)) in available_aliases.iter().enumerate() {
+            // Alias of the select item
             let col_name = column_binding.column_name.as_str();
-            // TODO(leiysky): check if expr is a qualified name
-            if let Expr::ColumnRef { column, .. } = expr {
-                if col_name == column.name.to_lowercase().as_str() {
-                    result.push(scalar.clone());
+            if let Expr::ColumnRef {
+                column,
+                database: None,
+                table: None,
+                ..
+            } = expr
+            {
+                if col_name.eq_ignore_ascii_case(column.name.as_str()) {
+                    result.push(i);
                 }
             }
         }
@@ -443,7 +460,32 @@ impl<'a> Binder {
                 format!("GROUP BY \"{}\" is ambiguous", expr),
             )))
         } else {
-            Ok((result[0].clone(), result[0].data_type()))
+            let (column_binding, scalar) = available_aliases[result[0]].clone();
+            // We will add the alias to BindContext, so we can reference it
+            // in `HAVING` and `ORDER BY` clause.
+            bind_context.columns.push(column_binding.clone());
+
+            let index = column_binding.index;
+            bind_context.aggregate_info.group_items.push(ScalarItem {
+                scalar: scalar.clone(),
+                index,
+            });
+            bind_context.aggregate_info.group_items_map.insert(
+                format!("{:?}", &scalar),
+                bind_context.aggregate_info.group_items.len() - 1,
+            );
+
+            // Add a mapping (alias -> scalar), so we can resolve the alias later
+            let column_ref: Scalar = BoundColumnRef {
+                column: column_binding,
+            }
+            .into();
+            bind_context.aggregate_info.group_items_map.insert(
+                format!("{:?}", &column_ref),
+                bind_context.aggregate_info.group_items.len() - 1,
+            );
+
+            Ok((scalar.clone(), scalar.data_type()))
         }
     }
 }

query/src/sql/planner/binder/select.rs

@@ -67,7 +67,12 @@ impl<'a> Binder {
 
         let (mut scalar_items, projections) = self.analyze_projection(&select_list)?;
 
+        // This will potentially add some alias group items to `from_context` if find some.
+        self.analyze_group_items(&mut from_context, &select_list, &stmt.group_by)
+            .await?;
+
         self.analyze_aggregate_select(&mut from_context, &mut select_list)?;
+
         let having = if let Some(having) = &stmt.having {
             Some(
                 self.analyze_aggregate_having(&mut from_context, having)
@@ -76,6 +81,7 @@ impl<'a> Binder {
         } else {
             None
         };
+
         let order_items = self.analyze_order_items(
             &from_context,
             &scalar_items,
@@ -88,9 +94,7 @@ impl<'a> Binder {
             || !stmt.group_by.is_empty()
             || stmt.having.is_some()
         {
-            s_expr = self
-                .bind_aggregate(&mut from_context, &select_list, &stmt.group_by, s_expr)
-                .await?;
+            s_expr = self.bind_aggregate(&mut from_context, s_expr).await?;
 
             if let Some(having) = having {
                 s_expr = self.bind_having(&from_context, having, s_expr).await?;

tests/suites/0_stateless/20+_others/20_0001_planner_v2.result

@@ -101,6 +101,9 @@ NULL	2	3
 0
 685
 0.685
+====Having alias====
+0
+1
 ====INNER_JOIN====
 1	1
 2	2

tests/suites/0_stateless/20+_others/20_0001_planner_v2.sql

@@ -106,6 +106,10 @@ select avg(number > 314) from numbers(1000);
 
 drop table t;
 
+select '====Having alias====';
+select number as a from numbers(1) group by a having a = 0;
+select number+1 as a from numbers(1) group by a having a = 1;
+
 -- Inner join
 select '====INNER_JOIN====';
 create table t(a int);