feat(cubesql): In subquery support (#7851)

Author: waralexrom

packages/cubejs-backend-native/Cargo.lock

@@ -10,12 +10,12 @@ homepage = "https://cube.dev"
 
 [dependencies]
 arc-swap = "1"
-datafusion = { git = 'https://github.com/cube-js/arrow-datafusion.git', rev = "28a07c390e7195dfd657c85118dee8cb73fc6bf7", default-features = false, features = ["regex_expressions", "unicode_expressions"] }
+datafusion = { git = 'https://github.com/cube-js/arrow-datafusion.git', rev = "a93bb9641d201b2b42ee92321f07e87bbd357d0e", default-features = false, features = ["regex_expressions", "unicode_expressions"] }
 anyhow = "1.0"
 thiserror = "1.0.50"
 cubeclient = { path = "../cubeclient" }
 pg-srv = { path = "../pg-srv" }
-sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "347f769500e3305f1920d8b38832f483d8795bd3" }
+sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "6a54d27d3b75a04b9f9cbe309a83078aa54b32fd" }
 lazy_static = "1.4.0"
 base64 = "0.13.0"
 tokio = { version = "^1.35", features = ["full", "rt", "tracing"] }

rust/cubesql/Cargo.lock

@@ -498,6 +498,7 @@ fn filter_push_down(
             subqueries,
             input,
             schema,
+            types,
         }) => {
             // TODO: Push Filter down Subquery
             issue_filter(
@@ -516,6 +517,7 @@ fn filter_push_down(
                         optimizer_config,
                     )?),
                     schema: schema.clone(),
+                    types: types.clone(),
                 }),
             )
         }

rust/cubesql/cubesql/Cargo.toml

@@ -339,6 +339,7 @@ fn limit_push_down(
             subqueries,
             input,
             schema,
+            types,
         }) => {
             // TODO: Pushing Limit down Subquery?
             issue_limit(
@@ -359,6 +360,7 @@ fn limit_push_down(
                         optimizer_config,
                     )?),
                     schema: schema.clone(),
+                    types: types.clone(),
                 }),
             )
         }

rust/cubesql/cubesql/src/compile/engine/df/optimizers/filter_push_down.rs

@@ -295,6 +295,7 @@ fn sort_push_down(
             subqueries,
             input,
             schema,
+            types,
         }) => {
             // TODO: Pushing Sort down Subquery?
             issue_sort(
@@ -306,6 +307,7 @@ fn sort_push_down(
                         .collect::<Result<_>>()?,
                     input: Arc::new(sort_push_down(optimizer, input, None, optimizer_config)?),
                     schema: schema.clone(),
+                    types: types.clone(),
                 }),
             )
         }

rust/cubesql/cubesql/src/compile/engine/df/optimizers/limit_push_down.rs

@@ -48,7 +48,12 @@ pub fn rewrite(expr: &Expr, map: &HashMap<Column, Option<Expr>>) -> Result<Optio
                 right: Box::new(right),
             })
         }
-        Expr::AnyExpr { left, op, right } => {
+        Expr::AnyExpr {
+            left,
+            op,
+            right,
+            all,
+        } => {
             let rewrites = match (rewrite(left, map)?, rewrite(right, map)?) {
                 (Some(left), Some(right)) => Some((left, right)),
                 _ => None,
@@ -57,6 +62,7 @@ pub fn rewrite(expr: &Expr, map: &HashMap<Column, Option<Expr>>) -> Result<Optio
                 left: Box::new(left),
                 op: op.clone(),
                 right: Box::new(right),
+                all: all.clone(),
             })
         }
         Expr::Like(Like {
@@ -310,6 +316,21 @@ pub fn rewrite(expr: &Expr, map: &HashMap<Column, Option<Expr>>) -> Result<Optio
         // As rewrites are used to push things down or up the plan, wildcards
         // might change the selection and should be marked as non-rewrittable
         Expr::Wildcard | Expr::QualifiedWildcard { .. } => None,
+        Expr::InSubquery {
+            expr,
+            subquery,
+            negated,
+        } => {
+            let rewrites = match (rewrite(expr, map)?, rewrite(subquery, map)?) {
+                (Some(expr), Some(subquery)) => Some((expr, subquery)),
+                _ => None,
+            };
+            rewrites.map(|(expr, subquery)| Expr::InSubquery {
+                expr: Box::new(expr),
+                subquery: Box::new(subquery),
+                negated: negated.clone(),
+            })
+        }
     })
 }
 

rust/cubesql/cubesql/src/compile/engine/df/optimizers/sort_push_down.rs

@@ -10,21 +10,22 @@ use crate::{
         rewrite::{
             analysis::LogicalPlanAnalysis, rewriter::Rewriter, AggregateFunctionExprDistinct,
             AggregateFunctionExprFun, AggregateSplit, AggregateUDFExprFun, AliasExprAlias,
-            AnyExprOp, BetweenExprNegated, BinaryExprOp, CastExprDataType, ChangeUserMemberValue,
-            ColumnExprColumn, CubeScanAliasToCube, CubeScanLimit, CubeScanOffset,
-            CubeScanUngrouped, CubeScanWrapped, DimensionName, EmptyRelationProduceOneRow,
-            FilterMemberMember, FilterMemberOp, FilterMemberValues, FilterOpOp, InListExprNegated,
-            JoinJoinConstraint, JoinJoinType, JoinLeftOn, JoinRightOn, LikeExprEscapeChar,
-            LikeExprLikeType, LikeExprNegated, LikeType, LimitFetch, LimitSkip, LiteralExprValue,
-            LiteralMemberRelation, LiteralMemberValue, LogicalPlanLanguage, MeasureName,
-            MemberErrorError, OrderAsc, OrderMember, OuterColumnExprColumn,
-            OuterColumnExprDataType, ProjectionAlias, ProjectionSplit, QueryParamIndex,
-            ScalarFunctionExprFun, ScalarUDFExprFun, ScalarVariableExprDataType,
-            ScalarVariableExprVariable, SegmentMemberMember, SortExprAsc, SortExprNullsFirst,
-            TableScanFetch, TableScanProjection, TableScanSourceTableName, TableScanTableName,
-            TableUDFExprFun, TimeDimensionDateRange, TimeDimensionGranularity, TimeDimensionName,
-            TryCastExprDataType, UnionAlias, WindowFunctionExprFun, WindowFunctionExprWindowFrame,
-            WrappedSelectAlias, WrappedSelectJoinJoinType, WrappedSelectLimit, WrappedSelectOffset,
+            AnyExprAll, AnyExprOp, BetweenExprNegated, BinaryExprOp, CastExprDataType,
+            ChangeUserMemberValue, ColumnExprColumn, CubeScanAliasToCube, CubeScanLimit,
+            CubeScanOffset, CubeScanUngrouped, CubeScanWrapped, DimensionName,
+            EmptyRelationProduceOneRow, FilterMemberMember, FilterMemberOp, FilterMemberValues,
+            FilterOpOp, InListExprNegated, InSubqueryNegated, JoinJoinConstraint, JoinJoinType,
+            JoinLeftOn, JoinRightOn, LikeExprEscapeChar, LikeExprLikeType, LikeExprNegated,
+            LikeType, LimitFetch, LimitSkip, LiteralExprValue, LiteralMemberRelation,
+            LiteralMemberValue, LogicalPlanLanguage, MeasureName, MemberErrorError, OrderAsc,
+            OrderMember, OuterColumnExprColumn, OuterColumnExprDataType, ProjectionAlias,
+            ProjectionSplit, QueryParamIndex, ScalarFunctionExprFun, ScalarUDFExprFun,
+            ScalarVariableExprDataType, ScalarVariableExprVariable, SegmentMemberMember,
+            SortExprAsc, SortExprNullsFirst, SubqueryTypes, TableScanFetch, TableScanProjection,
+            TableScanSourceTableName, TableScanTableName, TableUDFExprFun, TimeDimensionDateRange,
+            TimeDimensionGranularity, TimeDimensionName, TryCastExprDataType, UnionAlias,
+            WindowFunctionExprFun, WindowFunctionExprWindowFrame, WrappedSelectAlias,
+            WrappedSelectJoinJoinType, WrappedSelectLimit, WrappedSelectOffset,
             WrappedSelectSelectType, WrappedSelectType, WrappedSelectUngrouped,
         },
     },
@@ -221,11 +222,18 @@ impl LogicalPlanToLanguageConverter {
                     graph.add(LogicalPlanLanguage::LiteralExpr([value]))
                 }
             }
-            Expr::AnyExpr { left, op, right } => {
+            Expr::AnyExpr {
+                left,
+                op,
+                right,
+                all,
+            } => {
                 let left = Self::add_expr_replace_params(graph, left, query_params)?;
                 let op = add_expr_data_node!(graph, op, AnyExprOp);
                 let right = Self::add_expr_replace_params(graph, right, query_params)?;
-                graph.add(LogicalPlanLanguage::AnyExpr([left, op, right]))
+                let all = add_expr_data_node!(graph, all, AnyExprAll);
+
+                graph.add(LogicalPlanLanguage::AnyExpr([left, op, right, all]))
             }
             Expr::BinaryExpr { left, op, right } => {
                 let left = Self::add_expr_replace_params(graph, left, query_params)?;
@@ -414,6 +422,17 @@ impl LogicalPlanToLanguageConverter {
                 let negated = add_expr_data_node!(graph, negated, InListExprNegated);
                 graph.add(LogicalPlanLanguage::InListExpr([expr, list, negated]))
             }
+            Expr::InSubquery {
+                expr,
+                subquery,
+                negated,
+            } => {
+                let expr = Self::add_expr_replace_params(graph, expr, query_params)?;
+                let subquery = Self::add_expr_replace_params(graph, subquery, query_params)?;
+                let negated = add_expr_data_node!(graph, negated, InSubqueryNegated);
+
+                graph.add(LogicalPlanLanguage::InSubquery([expr, subquery, negated]))
+            }
             Expr::Wildcard => graph.add(LogicalPlanLanguage::WildcardExpr([])),
             Expr::GetIndexedField { expr, key } => {
                 let expr = Self::add_expr_replace_params(graph, expr, query_params)?;
@@ -548,8 +567,9 @@ impl LogicalPlanToLanguageConverter {
                     self.add_logical_plan_replace_params(node.input.as_ref(), query_params)?;
                 let subqueries =
                     add_plan_list_node!(self, node.subqueries, query_params, SubquerySubqueries);
+                let types = add_data_node!(self, node.types, SubqueryTypes);
                 self.graph
-                    .add(LogicalPlanLanguage::Subquery([input, subqueries]))
+                    .add(LogicalPlanLanguage::Subquery([input, subqueries, types]))
             }
             LogicalPlan::TableUDFs(node) => {
                 let expr =
@@ -818,7 +838,13 @@ pub fn node_to_expr(
             let left = Box::new(to_expr(params[0].clone())?);
             let op = match_data_node!(node_by_id, params[1], AnyExprOp);
             let right = Box::new(to_expr(params[2].clone())?);
-            Expr::AnyExpr { left, op, right }
+            let all = match_data_node!(node_by_id, params[3], AnyExprAll);
+            Expr::AnyExpr {
+                left,
+                op,
+                right,
+                all,
+            }
         }
         LogicalPlanLanguage::BinaryExpr(params) => {
             let left = Box::new(to_expr(params[0].clone())?);
@@ -1006,6 +1032,16 @@ pub fn node_to_expr(
                 "QueryParam can't be evaluated as an Expr node".to_string(),
             ));
         }
+        LogicalPlanLanguage::InSubquery(params) => {
+            let expr = Box::new(to_expr(params[0].clone())?);
+            let subquery = Box::new(to_expr(params[1].clone())?);
+            let negated = match_data_node!(node_by_id, params[2], InSubqueryNegated);
+            Expr::InSubquery {
+                expr,
+                subquery,
+                negated,
+            }
+        }
         x => panic!("Unexpected expression node: {:?}", x),
     })
 }
@@ -1179,8 +1215,9 @@ impl LanguageToLogicalPlanConverter {
                     .into_iter()
                     .map(|n| self.to_logical_plan(n))
                     .collect::<Result<Vec<_>, _>>()?;
+                let types = match_data_node!(node_by_id, params[2], SubqueryTypes);
                 LogicalPlanBuilder::from(input)
-                    .subquery(subqueries)?
+                    .subquery(subqueries, types)?
                     .build()?
             }
             LogicalPlanLanguage::TableUDFs(params) => {