From a151770cea7b6834e76814e23eab32fd3d4d618c Mon Sep 17 00:00:00 2001 From: Nikita Sivukhin Date: Mon, 27 Oct 2025 16:34:49 +0400 Subject: [PATCH] add minimal support of index_methods in the query planner in order to make integration tests work --- core/translate/index.rs | 326 +++++++++++++++++++++------------------- 1 file changed, 173 insertions(+), 153 deletions(-) diff --git a/core/translate/index.rs b/core/translate/index.rs index 52ee1c196..03f4b25c5 100644 --- a/core/translate/index.rs +++ b/core/translate/index.rs @@ -110,12 +110,30 @@ pub fn translate_create_index( crate::bail_parse_error!("Error: table '{tbl_name}' is not a b-tree table."); }; let columns = resolve_sorted_columns(&tbl, &columns)?; - let custom_module = using.is_some(); - if !with_clause.is_empty() && !custom_module { + if !with_clause.is_empty() && using.is_none() { crate::bail_parse_error!( "Error: additional parameters are allowed only for custom module indices: '{idx_name}' is not custom module index" ); } + + let mut module = None; + if let Some(using) = &using { + let index_modules = &resolver.symbol_table.index_methods; + let using = using.as_str(); + let index_module = index_modules.get(using); + if index_module.is_none() { + crate::bail_parse_error!("Error: unknown module name '{}'", using); + } + if let Some(index_module) = index_module { + let parameters = resolve_module_parameters(with_clause)?; + module = Some(index_module.attach(&IndexMethodConfiguration { + table_name: tbl.name.clone(), + index_name: idx_name.clone(), + columns: columns.clone(), + parameters: parameters.clone(), + })?); + } + } let idx = Arc::new(Index { name: idx_name.clone(), table_name: tbl.name.clone(), @@ -142,7 +160,7 @@ pub fn translate_create_index( // Allocate the necessary cursors: // // 1. sqlite_schema_cursor_id - sqlite_schema table - // 2. btree_cursor_id - new index btree + // 2. index_cursor_id - new index cursor // 3. table_cursor_id - table we are creating the index on // 4. sorter_cursor_id - sorter // 5. pseudo_cursor_id - pseudo table to store the sorted index values @@ -150,7 +168,7 @@ pub fn translate_create_index( let sqlite_schema_cursor_id = program.alloc_cursor_id(CursorType::BTreeTable(sqlite_table.clone())); let table_ref = program.table_reference_counter.next(); - let btree_cursor_id = program.alloc_cursor_id(CursorType::BTreeIndex(idx.clone())); + let index_cursor_id = program.alloc_cursor_id(CursorType::BTreeIndex(idx.clone())); let table_cursor_id = program.alloc_cursor_id_keyed( CursorKey::table(table_ref), CursorType::BTreeTable(tbl.clone()), @@ -204,163 +222,165 @@ pub fn translate_create_index( Some(sql), )?; - // determine the order of the columns in the index for the sorter - let order = idx.columns.iter().map(|c| c.order).collect(); - // open the sorter and the pseudo table - program.emit_insn(Insn::SorterOpen { - cursor_id: sorter_cursor_id, - columns: columns.len(), - order, - collations: idx.columns.iter().map(|c| c.collation).collect(), - }); - let content_reg = program.alloc_register(); - program.emit_insn(Insn::OpenPseudo { - cursor_id: pseudo_cursor_id, - content_reg, - num_fields: columns.len() + 1, - }); - - // open the table we are creating the index on for reading - program.emit_insn(Insn::OpenRead { - cursor_id: table_cursor_id, - root_page: tbl.root_page, - db: 0, - }); - - let loop_start_label = program.allocate_label(); - let loop_end_label = program.allocate_label(); - program.emit_insn(Insn::Rewind { - cursor_id: table_cursor_id, - pc_if_empty: loop_end_label, - }); - program.preassign_label_to_next_insn(loop_start_label); - - // Loop start: - // Collect index values into start_reg..rowid_reg - // emit MakeRecord (index key + rowid) into record_reg. - // - // Then insert the record into the sorter - let mut skip_row_label = None; - if let Some(where_clause) = where_clause { - let label = program.allocate_label(); - translate_condition_expr( - &mut program, - &table_references, - &where_clause, - ConditionMetadata { - jump_if_condition_is_true: false, - jump_target_when_false: label, - jump_target_when_true: BranchOffset::Placeholder, - jump_target_when_null: label, - }, - resolver, - )?; - skip_row_label = Some(label); - } - - let start_reg = program.alloc_registers(columns.len() + 1); - for (i, col) in columns.iter().enumerate() { - program.emit_column_or_rowid(table_cursor_id, col.pos_in_table, start_reg + i); - } - let rowid_reg = start_reg + columns.len(); - program.emit_insn(Insn::RowId { - cursor_id: table_cursor_id, - dest: rowid_reg, - }); - let record_reg = program.alloc_register(); - program.emit_insn(Insn::MakeRecord { - start_reg, - count: columns.len() + 1, - dest_reg: record_reg, - index_name: Some(idx_name.clone()), - affinity_str: None, - }); - program.emit_insn(Insn::SorterInsert { - cursor_id: sorter_cursor_id, - record_reg, - }); - - if let Some(skip_row_label) = skip_row_label { - program.resolve_label(skip_row_label, program.offset()); - } - program.emit_insn(Insn::Next { - cursor_id: table_cursor_id, - pc_if_next: loop_start_label, - }); - program.preassign_label_to_next_insn(loop_end_label); - - // Open the index btree we created for writing to insert the - // newly sorted index records. - program.emit_insn(Insn::OpenWrite { - cursor_id: btree_cursor_id, - root_page: RegisterOrLiteral::Register(root_page_reg), - db: 0, - }); - - let sorted_loop_start = program.allocate_label(); - let sorted_loop_end = program.allocate_label(); - - // Sort the index records in the sorter - program.emit_insn(Insn::SorterSort { - cursor_id: sorter_cursor_id, - pc_if_empty: sorted_loop_end, - }); - - let sorted_record_reg = program.alloc_register(); - - if unique { - // Since the records to be inserted are sorted, we can compare prev with current and if they are equal, - // we fall through to Halt with a unique constraint violation error. - let goto_label = program.allocate_label(); - let label_after_sorter_compare = program.allocate_label(); - program.resolve_label(goto_label, program.offset()); - program.emit_insn(Insn::Goto { - target_pc: label_after_sorter_compare, - }); - program.preassign_label_to_next_insn(sorted_loop_start); - program.emit_insn(Insn::SorterCompare { + if module.is_none() { + // determine the order of the columns in the index for the sorter + let order = idx.columns.iter().map(|c| c.order).collect(); + // open the sorter and the pseudo table + program.emit_insn(Insn::SorterOpen { cursor_id: sorter_cursor_id, - sorted_record_reg, - num_regs: columns.len(), - pc_when_nonequal: goto_label, + columns: columns.len(), + order, + collations: idx.columns.iter().map(|c| c.collation).collect(), }); - program.emit_insn(Insn::Halt { - err_code: SQLITE_CONSTRAINT_UNIQUE, - description: format_unique_violation_desc(tbl_name.as_str(), &idx), + let content_reg = program.alloc_register(); + program.emit_insn(Insn::OpenPseudo { + cursor_id: pseudo_cursor_id, + content_reg, + num_fields: columns.len() + 1, }); - program.preassign_label_to_next_insn(label_after_sorter_compare); - } else { - program.preassign_label_to_next_insn(sorted_loop_start); + + // open the table we are creating the index on for reading + program.emit_insn(Insn::OpenRead { + cursor_id: table_cursor_id, + root_page: tbl.root_page, + db: 0, + }); + + let loop_start_label = program.allocate_label(); + let loop_end_label = program.allocate_label(); + program.emit_insn(Insn::Rewind { + cursor_id: table_cursor_id, + pc_if_empty: loop_end_label, + }); + program.preassign_label_to_next_insn(loop_start_label); + + // Loop start: + // Collect index values into start_reg..rowid_reg + // emit MakeRecord (index key + rowid) into record_reg. + // + // Then insert the record into the sorter + let mut skip_row_label = None; + if let Some(where_clause) = where_clause { + let label = program.allocate_label(); + translate_condition_expr( + &mut program, + &table_references, + &where_clause, + ConditionMetadata { + jump_if_condition_is_true: false, + jump_target_when_false: label, + jump_target_when_true: BranchOffset::Placeholder, + jump_target_when_null: label, + }, + resolver, + )?; + skip_row_label = Some(label); + } + + let start_reg = program.alloc_registers(columns.len() + 1); + for (i, col) in columns.iter().enumerate() { + program.emit_column_or_rowid(table_cursor_id, col.pos_in_table, start_reg + i); + } + let rowid_reg = start_reg + columns.len(); + program.emit_insn(Insn::RowId { + cursor_id: table_cursor_id, + dest: rowid_reg, + }); + let record_reg = program.alloc_register(); + program.emit_insn(Insn::MakeRecord { + start_reg, + count: columns.len() + 1, + dest_reg: record_reg, + index_name: Some(idx_name.clone()), + affinity_str: None, + }); + program.emit_insn(Insn::SorterInsert { + cursor_id: sorter_cursor_id, + record_reg, + }); + + if let Some(skip_row_label) = skip_row_label { + program.resolve_label(skip_row_label, program.offset()); + } + program.emit_insn(Insn::Next { + cursor_id: table_cursor_id, + pc_if_next: loop_start_label, + }); + program.preassign_label_to_next_insn(loop_end_label); + + // Open the index btree we created for writing to insert the + // newly sorted index records. + program.emit_insn(Insn::OpenWrite { + cursor_id: index_cursor_id, + root_page: RegisterOrLiteral::Register(root_page_reg), + db: 0, + }); + + let sorted_loop_start = program.allocate_label(); + let sorted_loop_end = program.allocate_label(); + + // Sort the index records in the sorter + program.emit_insn(Insn::SorterSort { + cursor_id: sorter_cursor_id, + pc_if_empty: sorted_loop_end, + }); + + let sorted_record_reg = program.alloc_register(); + + if unique { + // Since the records to be inserted are sorted, we can compare prev with current and if they are equal, + // we fall through to Halt with a unique constraint violation error. + let goto_label = program.allocate_label(); + let label_after_sorter_compare = program.allocate_label(); + program.resolve_label(goto_label, program.offset()); + program.emit_insn(Insn::Goto { + target_pc: label_after_sorter_compare, + }); + program.preassign_label_to_next_insn(sorted_loop_start); + program.emit_insn(Insn::SorterCompare { + cursor_id: sorter_cursor_id, + sorted_record_reg, + num_regs: columns.len(), + pc_when_nonequal: goto_label, + }); + program.emit_insn(Insn::Halt { + err_code: SQLITE_CONSTRAINT_UNIQUE, + description: format_unique_violation_desc(tbl_name.as_str(), &idx), + }); + program.preassign_label_to_next_insn(label_after_sorter_compare); + } else { + program.preassign_label_to_next_insn(sorted_loop_start); + } + + program.emit_insn(Insn::SorterData { + pseudo_cursor: pseudo_cursor_id, + cursor_id: sorter_cursor_id, + dest_reg: sorted_record_reg, + }); + + // seek to the end of the index btree to position the cursor for appending + program.emit_insn(Insn::SeekEnd { + cursor_id: index_cursor_id, + }); + // insert new index record + program.emit_insn(Insn::IdxInsert { + cursor_id: index_cursor_id, + record_reg: sorted_record_reg, + unpacked_start: None, // TODO: optimize with these to avoid decoding record twice + unpacked_count: None, + flags: IdxInsertFlags::new().use_seek(false), + }); + program.emit_insn(Insn::SorterNext { + cursor_id: sorter_cursor_id, + pc_if_next: sorted_loop_start, + }); + program.preassign_label_to_next_insn(sorted_loop_end); } - program.emit_insn(Insn::SorterData { - pseudo_cursor: pseudo_cursor_id, - cursor_id: sorter_cursor_id, - dest_reg: sorted_record_reg, - }); - - // seek to the end of the index btree to position the cursor for appending - program.emit_insn(Insn::SeekEnd { - cursor_id: btree_cursor_id, - }); - // insert new index record - program.emit_insn(Insn::IdxInsert { - cursor_id: btree_cursor_id, - record_reg: sorted_record_reg, - unpacked_start: None, // TODO: optimize with these to avoid decoding record twice - unpacked_count: None, - flags: IdxInsertFlags::new().use_seek(false), - }); - program.emit_insn(Insn::SorterNext { - cursor_id: sorter_cursor_id, - pc_if_next: sorted_loop_start, - }); - program.preassign_label_to_next_insn(sorted_loop_end); - // End of the outer loop // // Keep schema table open to emit ParseSchema, close the other cursors. - program.close_cursors(&[sorter_cursor_id, table_cursor_id, btree_cursor_id]); + program.close_cursors(&[sorter_cursor_id, table_cursor_id, index_cursor_id]); program.emit_insn(Insn::SetCookie { db: 0,