simplify simple count detection

This commit is contained in:
pedrocarlo
2025-05-10 22:36:43 -03:00
parent 977d09fd36
commit 9f726dbe62
3 changed files with 43 additions and 47 deletions

View File

@@ -289,7 +289,7 @@ pub fn emit_query<'a>(
OperationMode::SELECT,
)?;
if plan.is_simple_count {
if plan.is_simple_count() {
emit_simple_count(program, t_ctx, plan)?;
return Ok(t_ctx.reg_result_cols_start.unwrap());
}

View File

@@ -308,8 +308,6 @@ 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 <tbl>`. This is set after the SelectPlan is create
pub is_simple_count: bool,
}
impl SelectPlan {
@@ -345,6 +343,48 @@ impl SelectPlan {
pub fn group_by_sorter_column_count(&self) -> usize {
self.agg_args_count() + self.group_by_col_count() + self.non_group_by_non_agg_column_count()
}
/// Reference: https://github.com/sqlite/sqlite/blob/5db695197b74580c777b37ab1b787531f15f7f9f/src/select.c#L8613
///
/// Checks to see if the query is of the format `SELECT count(*) FROM <tbl>`
pub fn is_simple_count(&self) -> bool {
if !self.where_clause.is_empty()
|| self.aggregates.len() != 1
|| matches!(self.query_type, SelectQueryType::Subquery { .. })
|| self.table_references.len() != 1
|| self.result_columns.len() != 1
|| self.group_by.is_some()
|| self.contains_constant_false_condition
// TODO: (pedrocarlo) maybe can optimize to use the count optmization with more columns
{
return false;
}
let table_ref = self.table_references.first().unwrap();
if !matches!(table_ref.table, crate::schema::Table::BTree(..)) {
return false;
}
let agg = self.aggregates.first().unwrap();
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 = &self.result_columns.get(0).unwrap().expr;
if *result_col_expr != count && *result_col_expr != count_star {
return false;
}
true
}
}
#[allow(dead_code)]

View File

@@ -106,7 +106,6 @@ 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();
@@ -388,8 +387,6 @@ 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))
}
@@ -489,47 +486,6 @@ fn estimate_num_labels(select: &SelectPlan) -> usize {
num_labels
}
/// Reference: https://github.com/sqlite/sqlite/blob/5db695197b74580c777b37ab1b787531f15f7f9f/src/select.c#L8613
///
/// Checks to see if the query is of the format `SELECT count(*) FROM <tbl>`
pub fn is_simple_count(plan: &SelectPlan) -> bool {
if !plan.where_clause.is_empty()
|| plan.aggregates.len() != 1
|| matches!(plan.query_type, SelectQueryType::Subquery { .. })
|| plan.table_references.len() != 1
|| plan.result_columns.len() != 1
|| plan.group_by.is_some()
// TODO: (pedrocarlo) maybe can optimize to use the count optmization with more columns
{
return false;
}
let table_ref = plan.table_references.first().unwrap();
if !matches!(table_ref.table, crate::schema::Table::BTree(..)) {
return false;
}
let agg = plan.aggregates.first().unwrap();
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
}
pub fn emit_simple_count<'a>(
program: &mut ProgramBuilder,
_t_ctx: &mut TranslateCtx<'a>,