fix: expression ownership in the temp table of subquery

KKould · Mar 16, 2024 · f02a5ee · f02a5ee
1 parent f8488f4
commit f02a5ee
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 14 deletions.
diff --git a/src/binder/expr.rs b/src/binder/expr.rs
@@ -110,30 +110,30 @@ impl<'a, T: Transaction> Binder<'a, T> {
             }),
             Expr::Subquery(subquery) => {
                 let (sub_query, column) = self.bind_subquery(subquery)?;
-                self.context.sub_query(SubQueryType::SubQuery(sub_query));
-
-                if self.context.is_step(&QueryBindStep::Where) {
-                    Ok(self.bind_temp_column(column))
+                let (expr, sub_query) = if self.context.is_step(&QueryBindStep::Where) {
+                    self.bind_temp_table(column, sub_query)?
                 } else {
-                    Ok(ScalarExpression::ColumnRef(column))
-                }
+                    (ScalarExpression::ColumnRef(column), sub_query)
+                };
+                self.context.sub_query(SubQueryType::SubQuery(sub_query));
+                Ok(expr)
             }
             Expr::InSubquery {
                 expr,
                 subquery,
                 negated,
             } => {
                 let (sub_query, column) = self.bind_subquery(subquery)?;
-                self.context
-                    .sub_query(SubQueryType::InSubQuery(*negated, sub_query));
 
                 if !self.context.is_step(&QueryBindStep::Where) {
                     return Err(DatabaseError::UnsupportedStmt(
                         "'IN (SUBQUERY)' can only appear in `WHERE`".to_string(),
                     ));
                 }
 
-                let alias_expr = self.bind_temp_column(column);
+                let (alias_expr, sub_query) = self.bind_temp_table(column, sub_query)?;
+                self.context
+                    .sub_query(SubQueryType::InSubQuery(*negated, sub_query));
 
                 Ok(ScalarExpression::Binary {
                     op: expression::BinaryOperator::Eq,
@@ -203,16 +203,22 @@ impl<'a, T: Transaction> Binder<'a, T> {
         }
     }
 
-    fn bind_temp_column(&mut self, column: ColumnRef) -> ScalarExpression {
+    fn bind_temp_table(
+        &mut self,
+        column: ColumnRef,
+        sub_query: LogicalPlan,
+    ) -> Result<(ScalarExpression, LogicalPlan), DatabaseError> {
         let mut alias_column = ColumnCatalog::clone(&column);
         alias_column.set_table_name(self.context.temp_table());
 
-        ScalarExpression::Alias {
+        let alias_expr = ScalarExpression::Alias {
             expr: Box::new(ScalarExpression::ColumnRef(column)),
             alias: AliasType::Expr(Box::new(ScalarExpression::ColumnRef(Arc::new(
                 alias_column,
             )))),
-        }
+        };
+        let alias_plan = self.bind_project(sub_query, vec![alias_expr.clone()])?;
+        Ok((alias_expr, alias_plan))
     }
 
     fn bind_subquery(

diff --git a/src/binder/select.rs b/src/binder/select.rs
@@ -510,7 +510,7 @@ impl<'a, T: Transaction> Binder<'a, T> {
         Ok(FilterOperator::build(having, children, true))
     }
 
-    fn bind_project(
+    pub(crate) fn bind_project(
         &mut self,
         children: LogicalPlan,
         select_list: Vec<ScalarExpression>,
@@ -722,7 +722,10 @@ impl<'a, T: Transaction> Binder<'a, T> {
             } => {
                 match op {
                     BinaryOperator::Eq => {
-                        match (left_expr.unpack_alias_ref(), right_expr.unpack_alias_ref()) {
+                        match (
+                            left_expr.unpack_alias_inner_ref().unpack_alias_ref(),
+                            right_expr.unpack_alias_inner_ref().unpack_alias_ref(),
+                        ) {
                             // example: foo = bar
                             (ScalarExpression::ColumnRef(l), ScalarExpression::ColumnRef(r)) => {
                                 // reorder left and right joins keys to pattern: (left, right)

diff --git a/src/expression/mod.rs b/src/expression/mod.rs
@@ -136,6 +136,18 @@ impl ScalarExpression {
         }
     }
 
+    pub fn unpack_alias_inner_ref(&self) -> &ScalarExpression {
+        if let ScalarExpression::Alias {
+            alias: AliasType::Expr(expr),
+            ..
+        } = self
+        {
+            expr.unpack_alias_inner_ref()
+        } else {
+            self
+        }
+    }
+
     pub fn try_reference(&mut self, output_exprs: &[ScalarExpression]) {
         let fn_output_column = |expr: &ScalarExpression| expr.unpack_alias_ref().output_column();
         let self_column = fn_output_column(self);