From b69debb8c0dec28a36566facc87b55719cad1553 Mon Sep 17 00:00:00 2001 From: pedrocarlo Date: Fri, 2 May 2025 22:03:04 -0300 Subject: [PATCH 1/7] create count opcode --- core/vdbe/execute.rs | 11 +++++++++++ core/vdbe/explain.rs | 13 +++++++++++++ core/vdbe/insn.rs | 11 +++++++++++ 3 files changed, 35 insertions(+) diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 173233bf4..6df66af17 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -4749,6 +4749,17 @@ pub fn op_affinity( Ok(InsnFunctionStepResult::Step) } +pub fn op_count( + program: &Program, + state: &mut ProgramState, + insn: &Insn, + pager: &Rc, + mv_store: Option<&Rc>, +) -> Result { + state.pc += 1; + Ok(InsnFunctionStepResult::Step) +} + fn exec_lower(reg: &OwnedValue) -> Option { match reg { OwnedValue::Text(t) => Some(OwnedValue::build_text(&t.as_str().to_lowercase())), diff --git a/core/vdbe/explain.rs b/core/vdbe/explain.rs index e2720dcb7..ff5892ddb 100644 --- a/core/vdbe/explain.rs +++ b/core/vdbe/explain.rs @@ -1448,6 +1448,19 @@ pub fn insn_to_str( .join(", ") ), ), + Insn::Count { + cursor_id, + target_reg, + exact, + } => ( + "Count", + *cursor_id as i32, + *target_reg as i32, + if *exact { 0 } else { 1 }, + OwnedValue::build_text(""), + 0, + "".to_string(), + ), }; format!( "{:<4} {:<17} {:<4} {:<4} {:<4} {:<13} {:<2} {}", diff --git a/core/vdbe/insn.rs b/core/vdbe/insn.rs index 8eea17e48..587695d30 100644 --- a/core/vdbe/insn.rs +++ b/core/vdbe/insn.rs @@ -857,6 +857,16 @@ pub enum Insn { count: NonZeroUsize, affinities: String, }, + + /// Store the number of entries (an integer value) in the table or index opened by cursor P1 in register P2. + /// + /// If P3==0, then an exact count is obtained, which involves visiting every btree page of the table. + /// But if P3 is non-zero, an estimate is returned based on the current cursor position. + Count { + cursor_id: CursorID, + target_reg: usize, + exact: bool, + }, } impl Insn { @@ -977,6 +987,7 @@ impl Insn { Insn::NotFound { .. } => execute::op_not_found, Insn::Affinity { .. } => execute::op_affinity, Insn::IdxDelete { .. } => execute::op_idx_delete, + Insn::Count { .. } => execute::op_count, } } } From 655ceeca458ea8cba34b6a0a42a62dbfdd74fc31 Mon Sep 17 00:00:00 2001 From: pedrocarlo Date: Sat, 3 May 2025 20:28:32 -0300 Subject: [PATCH 2/7] correct count implementation --- core/storage/btree.rs | 170 ++++++++++++++++++++++++++++++++++++++ core/translate/emitter.rs | 36 +++++--- core/translate/select.rs | 67 ++++++++++++++- core/vdbe/execute.rs | 18 ++++ 4 files changed, 277 insertions(+), 14 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 0523edcb8..92c79a384 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -200,6 +200,20 @@ enum ReadPayloadOverflow { }, } +struct CountInfo { + state: CountState, + /// Maintain count of the number of records in the btree. Used for the `Count` opcode + count: usize, +} + +/// State Machine of Count Operation +#[derive(Debug, Clone, Copy)] +enum CountState { + LoadPage, + MoveParent, + MoveChild, +} + #[derive(Clone, Debug)] pub enum BTreeKey<'a> { TableRowId((u64, Option<&'a ImmutableRecord>)), @@ -289,6 +303,7 @@ enum CursorState { Write(WriteInfo), Destroy(DestroyInfo), Delete(DeleteInfo), + Count(CountInfo), } impl CursorState { @@ -331,6 +346,20 @@ impl CursorState { _ => None, } } + + fn count_info(&self) -> Option<&CountInfo> { + match self { + CursorState::Count(x) => Some(x), + _ => None, + } + } + + fn mut_count_info(&mut self) -> Option<&mut CountInfo> { + match self { + CursorState::Count(x) => Some(x), + _ => None, + } + } } enum OverflowState { @@ -4228,6 +4257,147 @@ impl BTreeCursor { _ => false, } } + + /// Count the number of entries in the b-tree + /// + /// Only supposed to be used in the context of a simple Count Select Statement + pub fn count(&mut self) -> Result> { + if let Some(_mv_cursor) = &self.mv_cursor { + todo!("Implement count for mvcc"); + } + + if let CursorState::None = &self.state { + self.move_to_root(); + self.state = CursorState::Count(CountInfo { + state: CountState::LoadPage, + count: 0, + }); + } + + loop { + let count_state = self + .state + .count_info() + .expect("count info should have been set") + .state; + + match count_state { + CountState::LoadPage => { + let mem_page_rc = self.stack.top(); + return_if_locked_maybe_load!(self.pager, mem_page_rc); + let mem_page = mem_page_rc.get(); + let contents = mem_page.contents.as_ref().unwrap(); + let cell_idx = self.stack.current_cell_index() as usize; + + let count_info = self + .state + .mut_count_info() + .expect("count info should have been set"); + + /* If this is a leaf page or the tree is not an int-key tree, then + ** this page contains countable entries. Increment the entry counter + ** accordingly. + */ + // TODO: (pedrocarlo) does not take into account if "tree is not an int-key tree" condition + if contents.is_leaf() { + count_info.count += contents.cell_count(); + } + + if contents.is_leaf() || cell_idx > contents.cell_count() { + count_info.state = CountState::MoveParent; + } else { + count_info.state = CountState::MoveChild; + } + } + CountState::MoveParent => { + let mem_page_rc = self.stack.top(); + assert!(mem_page_rc.is_loaded()); + + if !self.stack.has_parent() { + // All pages of the b-tree have been visited. Return successfully + self.move_to_root(); + // Borrow count_info here like this due to borrow checking + let count_info = self + .state + .mut_count_info() + .expect("count info should have been set"); + return Ok(CursorResult::Ok(count_info.count)); + } + + let count_info = self + .state + .mut_count_info() + .expect("count info should have been set"); + + // Move to parent + self.going_upwards = true; + self.stack.pop(); + + count_info.state = CountState::LoadPage; + } + CountState::MoveChild => { + let mem_page_rc = self.stack.top(); + assert!(mem_page_rc.is_loaded()); + let contents = mem_page_rc.get().contents.as_ref().unwrap(); + let cell_idx = self.stack.current_cell_index() as usize; + + assert!(!contents.is_leaf()); + assert!( + cell_idx <= contents.cell_count(), + "cell_idx={} cell_count={}", + cell_idx, + contents.cell_count() + ); + + if cell_idx == contents.cell_count() { + // Move to right child + // should be safe as contents is not a leaf page + let right_most_pointer = contents.rightmost_pointer().unwrap(); + self.stack.advance(); + let mem_page = self.pager.read_page(right_most_pointer as usize)?; + self.going_upwards = false; + self.stack.push(mem_page); + } else { + // Move to child left page + let cell = contents.cell_get( + cell_idx, + payload_overflow_threshold_max( + contents.page_type(), + self.usable_space() as u16, + ), + payload_overflow_threshold_min( + contents.page_type(), + self.usable_space() as u16, + ), + self.usable_space(), + )?; + + match cell { + BTreeCell::TableInteriorCell(TableInteriorCell { + _left_child_page: left_child_page, + .. + }) + | BTreeCell::IndexInteriorCell(IndexInteriorCell { + left_child_page, + .. + }) => { + self.stack.advance(); + let mem_page = self.pager.read_page(left_child_page as usize)?; + self.going_upwards = false; + self.stack.push(mem_page); + } + _ => unreachable!(), + } + } + let count_info = self + .state + .mut_count_info() + .expect("count info should have been set"); + count_info.state = CountState::LoadPage; + } + } + } + } } #[cfg(debug_assertions)] diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index 4dc72249b..5dc137207 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -23,6 +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::schema::ParseSchema; use super::subquery::emit_subqueries; @@ -232,6 +233,8 @@ 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)?; @@ -309,20 +312,23 @@ pub fn emit_query<'a>( program.preassign_label_to_next_insn(jump_target_when_true); } - // Set up main query execution loop - open_loop( - program, - t_ctx, - &plan.table_references, - &plan.join_order, - &plan.where_clause, - )?; + if !is_simple_count { + // Set up main query execution loop + open_loop( + program, + t_ctx, + &plan.table_references, + &plan.join_order, + &plan.where_clause, + )?; - // Process result columns and expressions in the inner loop - emit_loop(program, t_ctx, plan)?; + // Process result columns and expressions in the inner loop + emit_loop(program, t_ctx, plan)?; + + // Clean up and close the main execution loop + close_loop(program, t_ctx, &plan.table_references, &plan.join_order)?; + } - // Clean up and close the main execution loop - close_loop(program, t_ctx, &plan.table_references, &plan.join_order)?; program.preassign_label_to_next_insn(after_main_loop_label); let mut order_by_necessary = plan.order_by.is_some() && !plan.contains_constant_false_condition; @@ -341,7 +347,11 @@ 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 - emit_ungrouped_aggregation(program, t_ctx, plan)?; + if is_simple_count { + emit_simple_count(program, t_ctx, plan)?; + } else { + emit_ungrouped_aggregation(program, t_ctx, plan)?; + } // Single row result for aggregates without GROUP BY, so ORDER BY not needed order_by_necessary = false; } diff --git a/core/translate/select.rs b/core/translate/select.rs index 1bf6905da..d54e0a907 100644 --- a/core/translate/select.rs +++ b/core/translate/select.rs @@ -1,4 +1,4 @@ -use super::emitter::emit_program; +use super::emitter::{emit_program, TranslateCtx}; use super::plan::{select_star, JoinOrderMember, Operation, Search, SelectQueryType}; use super::planner::Scope; use crate::function::{AggFunc, ExtFunc, Func}; @@ -10,6 +10,7 @@ use crate::translate::planner::{ }; use crate::util::normalize_ident; use crate::vdbe::builder::{ProgramBuilderOpts, QueryMode}; +use crate::vdbe::insn::Insn; use crate::SymbolTable; use crate::{schema::Schema, vdbe::builder::ProgramBuilder, Result}; use limbo_sqlite3_parser::ast::{self, SortOrder}; @@ -484,3 +485,67 @@ 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 ` +pub fn is_simple_count(plan: &SelectPlan) -> bool { + // TODO: (pedrocarlo) check for HAVING clause + 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; + } + true +} + +pub fn emit_simple_count<'a>( + program: &mut ProgramBuilder, + _t_ctx: &mut TranslateCtx<'a>, + plan: &'a SelectPlan, +) -> Result<()> { + let cursors = plan + .table_references + .get(0) + .unwrap() + .resolve_cursors(program)?; + + let cursor_id = { + match cursors { + (_, Some(cursor_id)) | (Some(cursor_id), None) => cursor_id, + _ => panic!("cursor for table should have been opened"), + } + }; + + // TODO: I think this allocation can be avoided if we are smart with the `TranslateCtx` + let target_reg = program.alloc_register(); + + program.emit_insn(Insn::Count { + cursor_id, + target_reg, + exact: true, + }); + + program.emit_insn(Insn::Close { cursor_id }); + let output_reg = program.alloc_register(); + program.emit_insn(Insn::Copy { + src_reg: target_reg, + dst_reg: output_reg, + amount: 0, + }); + program.emit_result_row(output_reg, 1); + Ok(()) +} diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 6df66af17..d6ab6ba4d 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -4756,6 +4756,24 @@ pub fn op_count( pager: &Rc, mv_store: Option<&Rc>, ) -> Result { + let Insn::Count { + cursor_id, + target_reg, + exact, + } = insn + else { + unreachable!("unexpected Insn {:?}", insn) + }; + + let count = { + let mut cursor = must_be_btree_cursor!(*cursor_id, program.cursor_ref, state, "Count"); + let cursor = cursor.as_btree_mut(); + let count = return_if_io!(cursor.count()); + count + }; + + state.registers[*target_reg] = Register::OwnedValue(OwnedValue::Integer(count as i64)); + state.pc += 1; Ok(InsnFunctionStepResult::Step) } From 342bf51c88b85daed3f16de7d06889a5cef0b8a8 Mon Sep 17 00:00:00 2001 From: pedrocarlo Date: Sat, 3 May 2025 23:11:35 -0300 Subject: [PATCH 3/7] remove state machine for count --- core/storage/btree.rs | 212 +++++++++++++++--------------------------- 1 file changed, 76 insertions(+), 136 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 92c79a384..1e6175473 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -200,20 +200,6 @@ enum ReadPayloadOverflow { }, } -struct CountInfo { - state: CountState, - /// Maintain count of the number of records in the btree. Used for the `Count` opcode - count: usize, -} - -/// State Machine of Count Operation -#[derive(Debug, Clone, Copy)] -enum CountState { - LoadPage, - MoveParent, - MoveChild, -} - #[derive(Clone, Debug)] pub enum BTreeKey<'a> { TableRowId((u64, Option<&'a ImmutableRecord>)), @@ -303,7 +289,6 @@ enum CursorState { Write(WriteInfo), Destroy(DestroyInfo), Delete(DeleteInfo), - Count(CountInfo), } impl CursorState { @@ -346,20 +331,6 @@ impl CursorState { _ => None, } } - - fn count_info(&self) -> Option<&CountInfo> { - match self { - CursorState::Count(x) => Some(x), - _ => None, - } - } - - fn mut_count_info(&mut self) -> Option<&mut CountInfo> { - match self { - CursorState::Count(x) => Some(x), - _ => None, - } - } } enum OverflowState { @@ -394,6 +365,8 @@ pub struct BTreeCursor { reusable_immutable_record: RefCell>, empty_record: Cell, pub index_key_sort_order: IndexKeySortOrder, + /// Maintain count of the number of records in the btree. Used for the `Count` opcode + count: usize, } impl BTreeCursor { @@ -419,6 +392,7 @@ impl BTreeCursor { reusable_immutable_record: RefCell::new(None), empty_record: Cell::new(true), index_key_sort_order: IndexKeySortOrder::default(), + count: 0, } } @@ -4262,138 +4236,104 @@ impl BTreeCursor { /// /// Only supposed to be used in the context of a simple Count Select Statement pub fn count(&mut self) -> Result> { + if self.count == 0 { + self.move_to_root(); + } + if let Some(_mv_cursor) = &self.mv_cursor { todo!("Implement count for mvcc"); } - if let CursorState::None = &self.state { - self.move_to_root(); - self.state = CursorState::Count(CountInfo { - state: CountState::LoadPage, - count: 0, - }); - } + let mut mem_page_rc; + let mut mem_page; + let mut contents; loop { - let count_state = self - .state - .count_info() - .expect("count info should have been set") - .state; + mem_page_rc = self.stack.top(); + return_if_locked_maybe_load!(self.pager, mem_page_rc); + mem_page = mem_page_rc.get(); + contents = mem_page.contents.as_ref().unwrap(); - match count_state { - CountState::LoadPage => { - let mem_page_rc = self.stack.top(); - return_if_locked_maybe_load!(self.pager, mem_page_rc); - let mem_page = mem_page_rc.get(); - let contents = mem_page.contents.as_ref().unwrap(); - let cell_idx = self.stack.current_cell_index() as usize; + /* If this is a leaf page or the tree is not an int-key tree, then + ** this page contains countable entries. Increment the entry counter + ** accordingly. + */ + // TODO: (pedrocarlo) does not take into account if "tree is not an int-key tree" condition + if contents.is_leaf() { + self.count += contents.cell_count(); + } - let count_info = self - .state - .mut_count_info() - .expect("count info should have been set"); - - /* If this is a leaf page or the tree is not an int-key tree, then - ** this page contains countable entries. Increment the entry counter - ** accordingly. - */ - // TODO: (pedrocarlo) does not take into account if "tree is not an int-key tree" condition - if contents.is_leaf() { - count_info.count += contents.cell_count(); - } - - if contents.is_leaf() || cell_idx > contents.cell_count() { - count_info.state = CountState::MoveParent; - } else { - count_info.state = CountState::MoveChild; - } - } - CountState::MoveParent => { - let mem_page_rc = self.stack.top(); - assert!(mem_page_rc.is_loaded()); + let cell_idx = self.stack.current_cell_index() as usize; + // Second condition is necessary in case we return if the page is locked in the loop below + if contents.is_leaf() || cell_idx > contents.cell_count() { + loop { if !self.stack.has_parent() { // All pages of the b-tree have been visited. Return successfully self.move_to_root(); - // Borrow count_info here like this due to borrow checking - let count_info = self - .state - .mut_count_info() - .expect("count info should have been set"); - return Ok(CursorResult::Ok(count_info.count)); - } - let count_info = self - .state - .mut_count_info() - .expect("count info should have been set"); + return Ok(CursorResult::Ok(self.count)); + } // Move to parent self.going_upwards = true; self.stack.pop(); - count_info.state = CountState::LoadPage; - } - CountState::MoveChild => { - let mem_page_rc = self.stack.top(); - assert!(mem_page_rc.is_loaded()); - let contents = mem_page_rc.get().contents.as_ref().unwrap(); + mem_page_rc = self.stack.top(); + return_if_locked_maybe_load!(self.pager, mem_page_rc); + mem_page = mem_page_rc.get(); + contents = mem_page.contents.as_ref().unwrap(); + let cell_idx = self.stack.current_cell_index() as usize; - assert!(!contents.is_leaf()); - assert!( - cell_idx <= contents.cell_count(), - "cell_idx={} cell_count={}", - cell_idx, - contents.cell_count() - ); + if !(cell_idx > contents.cell_count()) { + break; + } + } + } - if cell_idx == contents.cell_count() { - // Move to right child - // should be safe as contents is not a leaf page - let right_most_pointer = contents.rightmost_pointer().unwrap(); + let cell_idx = self.stack.current_cell_index() as usize; + + assert!(cell_idx <= contents.cell_count(),); + assert!(!contents.is_leaf()); + + if cell_idx == contents.cell_count() { + // Move to right child + // should be safe as contents is not a leaf page + let right_most_pointer = contents.rightmost_pointer().unwrap(); + self.stack.advance(); + let mem_page = self.pager.read_page(right_most_pointer as usize)?; + self.going_upwards = false; + self.stack.push(mem_page); + } else { + // Move to child left page + let cell = contents.cell_get( + cell_idx, + payload_overflow_threshold_max( + contents.page_type(), + self.usable_space() as u16, + ), + payload_overflow_threshold_min( + contents.page_type(), + self.usable_space() as u16, + ), + self.usable_space(), + )?; + + match cell { + BTreeCell::TableInteriorCell(TableInteriorCell { + _left_child_page: left_child_page, + .. + }) + | BTreeCell::IndexInteriorCell(IndexInteriorCell { + left_child_page, .. + }) => { self.stack.advance(); - let mem_page = self.pager.read_page(right_most_pointer as usize)?; + let mem_page = self.pager.read_page(left_child_page as usize)?; self.going_upwards = false; self.stack.push(mem_page); - } else { - // Move to child left page - let cell = contents.cell_get( - cell_idx, - payload_overflow_threshold_max( - contents.page_type(), - self.usable_space() as u16, - ), - payload_overflow_threshold_min( - contents.page_type(), - self.usable_space() as u16, - ), - self.usable_space(), - )?; - - match cell { - BTreeCell::TableInteriorCell(TableInteriorCell { - _left_child_page: left_child_page, - .. - }) - | BTreeCell::IndexInteriorCell(IndexInteriorCell { - left_child_page, - .. - }) => { - self.stack.advance(); - let mem_page = self.pager.read_page(left_child_page as usize)?; - self.going_upwards = false; - self.stack.push(mem_page); - } - _ => unreachable!(), - } } - let count_info = self - .state - .mut_count_info() - .expect("count info should have been set"); - count_info.state = CountState::LoadPage; + _ => unreachable!(), } } } From e9b1631d3c86d5afd86bbd3452e9d2fbb68a91f3 Mon Sep 17 00:00:00 2001 From: pedrocarlo Date: Sat, 3 May 2025 23:44:20 -0300 Subject: [PATCH 4/7] fix is_simple_count detection --- core/translate/emitter.rs | 8 +++----- core/translate/plan.rs | 2 ++ core/translate/select.rs | 19 +++++++++++++++++++ 3 files changed, 24 insertions(+), 5 deletions(-) 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 } From 7508043b62d2d4500fce223cca949ad21e34e4a8 Mon Sep 17 00:00:00 2001 From: pedrocarlo Date: Sun, 4 May 2025 01:32:09 -0300 Subject: [PATCH 5/7] add bench for select count --- core/benches/benchmark.rs | 55 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 54 insertions(+), 1 deletion(-) diff --git a/core/benches/benchmark.rs b/core/benches/benchmark.rs index 1fed5308b..388fcf933 100644 --- a/core/benches/benchmark.rs +++ b/core/benches/benchmark.rs @@ -180,9 +180,62 @@ fn bench_execute_select_1(criterion: &mut Criterion) { group.finish(); } +fn bench_execute_select_count(criterion: &mut Criterion) { + // https://github.com/tursodatabase/limbo/issues/174 + // The rusqlite benchmark crashes on Mac M1 when using the flamegraph features + let enable_rusqlite = std::env::var("DISABLE_RUSQLITE_BENCHMARK").is_err(); + + #[allow(clippy::arc_with_non_send_sync)] + let io = Arc::new(PlatformIO::new().unwrap()); + let db = Database::open_file(io.clone(), "../testing/testing.db", false).unwrap(); + let limbo_conn = db.connect().unwrap(); + + let mut group = criterion.benchmark_group("Execute `SELECT count() FROM users`"); + + group.bench_function("limbo_execute_select_count", |b| { + let mut stmt = limbo_conn.prepare("SELECT count() FROM users").unwrap(); + let io = io.clone(); + b.iter(|| { + loop { + match stmt.step().unwrap() { + limbo_core::StepResult::Row => { + black_box(stmt.row()); + } + limbo_core::StepResult::IO => { + let _ = io.run_once(); + } + limbo_core::StepResult::Done => { + break; + } + limbo_core::StepResult::Interrupt | limbo_core::StepResult::Busy => { + unreachable!(); + } + } + } + stmt.reset(); + }); + }); + + if enable_rusqlite { + let sqlite_conn = rusqlite_open(); + + group.bench_function("sqlite_execute_select_count", |b| { + let mut stmt = sqlite_conn.prepare("SELECT count() FROM users").unwrap(); + b.iter(|| { + let mut rows = stmt.raw_query(); + while let Some(row) = rows.next().unwrap() { + black_box(row); + } + }); + }); + } + + group.finish(); +} + criterion_group! { name = benches; config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None))); - targets = bench_prepare_query, bench_execute_select_1, bench_execute_select_rows + targets = bench_prepare_query, bench_execute_select_1, bench_execute_select_rows, bench_execute_select_count } criterion_main!(benches); From 977d09fd3645fce55c376d212f7163b985c218a7 Mon Sep 17 00:00:00 2001 From: pedrocarlo Date: Thu, 8 May 2025 17:17:20 -0300 Subject: [PATCH 6/7] small fixes --- core/storage/btree.rs | 5 ++--- core/translate/emitter.rs | 39 +++++++++++++++++++-------------------- core/translate/select.rs | 3 +-- 3 files changed, 22 insertions(+), 25 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 1e6175473..208900951 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -4258,8 +4258,7 @@ impl BTreeCursor { ** this page contains countable entries. Increment the entry counter ** accordingly. */ - // TODO: (pedrocarlo) does not take into account if "tree is not an int-key tree" condition - if contents.is_leaf() { + if !matches!(contents.page_type(), PageType::TableInterior) { self.count += contents.cell_count(); } @@ -4286,7 +4285,7 @@ impl BTreeCursor { let cell_idx = self.stack.current_cell_index() as usize; - if !(cell_idx > contents.cell_count()) { + if cell_idx <= contents.cell_count() { break; } } diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index 08ae6b7e8..65129a9f0 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -23,8 +23,8 @@ 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; use super::schema::ParseSchema; +use super::select::emit_simple_count; use super::subquery::emit_subqueries; #[derive(Debug)] @@ -289,6 +289,11 @@ pub fn emit_query<'a>( OperationMode::SELECT, )?; + if plan.is_simple_count { + emit_simple_count(program, t_ctx, plan)?; + return Ok(t_ctx.reg_result_cols_start.unwrap()); + } + for where_term in plan .where_clause .iter() @@ -310,22 +315,20 @@ pub fn emit_query<'a>( program.preassign_label_to_next_insn(jump_target_when_true); } - if !plan.is_simple_count { - // Set up main query execution loop - open_loop( - program, - t_ctx, - &plan.table_references, - &plan.join_order, - &plan.where_clause, - )?; + // Set up main query execution loop + open_loop( + program, + t_ctx, + &plan.table_references, + &plan.join_order, + &plan.where_clause, + )?; - // Process result columns and expressions in the inner loop - emit_loop(program, t_ctx, plan)?; + // Process result columns and expressions in the inner loop + emit_loop(program, t_ctx, plan)?; - // Clean up and close the main execution loop - close_loop(program, t_ctx, &plan.table_references, &plan.join_order)?; - } + // Clean up and close the main execution loop + close_loop(program, t_ctx, &plan.table_references, &plan.join_order)?; program.preassign_label_to_next_insn(after_main_loop_label); @@ -345,11 +348,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 plan.is_simple_count { - emit_simple_count(program, t_ctx, plan)?; - } else { - emit_ungrouped_aggregation(program, t_ctx, plan)?; - } + emit_ungrouped_aggregation(program, t_ctx, plan)?; // Single row result for aggregates without GROUP BY, so ORDER BY not needed order_by_necessary = false; } diff --git a/core/translate/select.rs b/core/translate/select.rs index 4f39b827f..90819169a 100644 --- a/core/translate/select.rs +++ b/core/translate/select.rs @@ -493,7 +493,6 @@ fn estimate_num_labels(select: &SelectPlan) -> usize { /// /// Checks to see if the query is of the format `SELECT count(*) FROM ` pub fn is_simple_count(plan: &SelectPlan) -> bool { - // TODO: (pedrocarlo) check for HAVING clause if !plan.where_clause.is_empty() || plan.aggregates.len() != 1 || matches!(plan.query_type, SelectQueryType::Subquery { .. }) @@ -525,7 +524,7 @@ pub fn is_simple_count(plan: &SelectPlan) -> bool { filter_over: None, }; let result_col_expr = &plan.result_columns.get(0).unwrap().expr; - if *result_col_expr != count || *result_col_expr != count_star { + if *result_col_expr != count && *result_col_expr != count_star { return false; } true From 9f726dbe62298f35bb0b96f30a2faf3559d8b050 Mon Sep 17 00:00:00 2001 From: pedrocarlo Date: Sat, 10 May 2025 22:36:43 -0300 Subject: [PATCH 7/7] simplify simple count detection --- core/translate/emitter.rs | 2 +- core/translate/plan.rs | 44 +++++++++++++++++++++++++++++++++++++-- core/translate/select.rs | 44 --------------------------------------- 3 files changed, 43 insertions(+), 47 deletions(-) diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index 65129a9f0..6557ca169 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -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()); } diff --git a/core/translate/plan.rs b/core/translate/plan.rs index 0ee865cb5..e34efb29b 100644 --- a/core/translate/plan.rs +++ b/core/translate/plan.rs @@ -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 `. 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 ` + 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)] diff --git a/core/translate/select.rs b/core/translate/select.rs index 90819169a..d6766a20f 100644 --- a/core/translate/select.rs +++ b/core/translate/select.rs @@ -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 ` -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>,