diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index 5dc137207..08ae6b7e8 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -23,7 +23,7 @@ use super::group_by::{ use super::main_loop::{close_loop, emit_loop, init_loop, open_loop, LeftJoinMetadata, LoopLabels}; use super::order_by::{emit_order_by, init_order_by, SortMetadata}; use super::plan::{JoinOrderMember, Operation, SelectPlan, TableReference, UpdatePlan}; -use super::select::{emit_simple_count, is_simple_count}; +use super::select::emit_simple_count; use super::schema::ParseSchema; use super::subquery::emit_subqueries; @@ -233,8 +233,6 @@ pub fn emit_query<'a>( plan: &'a mut SelectPlan, t_ctx: &'a mut TranslateCtx<'a>, ) -> Result { - let is_simple_count = is_simple_count(plan); - // Emit subqueries first so the results can be read in the main query loop. emit_subqueries(program, t_ctx, &mut plan.table_references)?; @@ -312,7 +310,7 @@ pub fn emit_query<'a>( program.preassign_label_to_next_insn(jump_target_when_true); } - if !is_simple_count { + if !plan.is_simple_count { // Set up main query execution loop open_loop( program, @@ -347,7 +345,7 @@ pub fn emit_query<'a>( group_by_emit_row_phase(program, t_ctx, plan)?; } else if !plan.aggregates.is_empty() { // Handle aggregation without GROUP BY - if is_simple_count { + if plan.is_simple_count { emit_simple_count(program, t_ctx, plan)?; } else { emit_ungrouped_aggregation(program, t_ctx, plan)?; diff --git a/core/translate/plan.rs b/core/translate/plan.rs index dd73e7e26..0ee865cb5 100644 --- a/core/translate/plan.rs +++ b/core/translate/plan.rs @@ -308,6 +308,8 @@ pub struct SelectPlan { pub contains_constant_false_condition: bool, /// query type (top level or subquery) pub query_type: SelectQueryType, + /// if the query is of the format `SELECT count(*) FROM `. This is set after the SelectPlan is create + pub is_simple_count: bool, } impl SelectPlan { diff --git a/core/translate/select.rs b/core/translate/select.rs index d54e0a907..4f39b827f 100644 --- a/core/translate/select.rs +++ b/core/translate/select.rs @@ -106,6 +106,7 @@ pub fn prepare_select_plan<'a>( offset: None, contains_constant_false_condition: false, query_type: SelectQueryType::TopLevel, + is_simple_count: false, }; let mut aggregate_expressions = Vec::new(); @@ -387,6 +388,8 @@ pub fn prepare_select_plan<'a>( (plan.limit, plan.offset) = select.limit.map_or(Ok((None, None)), |l| parse_limit(&l))?; + plan.is_simple_count = is_simple_count(&plan); + // Return the unoptimized query plan Ok(Plan::Select(plan)) } @@ -509,6 +512,22 @@ pub fn is_simple_count(plan: &SelectPlan) -> bool { if !matches!(agg.func, AggFunc::Count0) { return false; } + + let count = limbo_sqlite3_parser::ast::Expr::FunctionCall { + name: limbo_sqlite3_parser::ast::Id("count".to_string()), + distinctness: None, + args: None, + order_by: None, + filter_over: None, + }; + let count_star = limbo_sqlite3_parser::ast::Expr::FunctionCallStar { + name: limbo_sqlite3_parser::ast::Id("count".to_string()), + filter_over: None, + }; + let result_col_expr = &plan.result_columns.get(0).unwrap().expr; + if *result_col_expr != count || *result_col_expr != count_star { + return false; + } true }