优化和收尾

chusql-storage/src/heap.rs

@@ -1,3 +1,4 @@
+use std::collections::HashSet;
 use std::io;
 use std::path::Path;
 
@@ -6,7 +7,10 @@ use crate::page::PageId;
 use crate::page::{Page, PageFile};
 use crate::protocol::Row;
 
-// 堆表：行按插入顺序进页，支持全表扫描与按索引取行。
+// 堆表：行按插入顺序进页，支持全表扫描、按索引取行、行级删除。
+//
+// 索引由**表自己维护**：有哪些索引记在数据字典里，插入时按索引的列从行里取值写进去，
+// 调用方不需要知道有哪些索引，也不需要自己指定 key。
 
 // 页内布局
 /// 页头：slot_count 占 2 字节
@@ -47,6 +51,14 @@ fn unpack_position(v: u64) -> (PageId, u16) {
     (v >> 16, (v & 0xFFFF) as u16)
 }
 
+/// 取行里某一列的整数值（不是整数就没有）
+fn row_int(row: &Row, column: &str) -> Option<i64> {
+    match row.get(column) {
+        Some(serde_json::Value::Number(n)) => n.as_i64(),
+        _ => None,
+    }
+}
+
 // 页内操作
 /// 只管一页里的字节
 pub struct HeapPage;
@@ -103,6 +115,20 @@ impl HeapPage {
         Some(n)
     }
 
+    /// 删掉一个槽位：只把长度置 0，行字节原地留着（读的时候按"长度 0 = 没有这一行"处理）
+    pub fn delete_tuple(page: &mut Page, slot: u16) -> bool {
+        let n = Self::slot_count(page);
+        if slot >= n {
+            return false;
+        }
+        let slot_off = SLOT_DIR_START + (slot as usize) * SLOT_SIZE;
+        if read_u32(&page.data, slot_off + 4) == 0 {
+            return false;
+        }
+        write_u32(&mut page.data, slot_off + 4, 0);
+        true
+    }
+
     /// 读槽位里的行字节
     pub fn get_tuple(page: &Page, slot: u16) -> Option<Vec<u8>> {
         let n = Self::slot_count(page);
@@ -124,27 +150,30 @@ impl HeapPage {
     }
 }
 
+// 索引
+/// 表上的一个索引：盯着一列，键是这一列的整数值
+struct NamedIndex {
+    column: String,
+    tree: DiskBTree,
+}
+
 // 堆表
-/// 一个文件 = 若干页
+/// 一个数据文件 = 若干页；外加若干索引文件
 pub struct HeapTable {
     file: PageFile,
-    index: Option<DiskBTree>,
+    indexes: Vec<NamedIndex>,
 }
 
 impl HeapTable {
     /// 打开表（不带索引）
-    pub fn open<P: AsRef<Path>>(
-        path: P,
-        page_size: usize,
-        pool_size: usize,
-    ) -> io::Result<Self> {
-         Ok(HeapTable {
+    pub fn open<P: AsRef<Path>>(path: P, page_size: usize, pool_size: usize) -> io::Result<Self> {
+        Ok(HeapTable {
             file: PageFile::with_options(path, page_size, pool_size)?,
-            index: None,
+            indexes: Vec::new(),
         })
     }
 
-    /// 打开表 + 索引
+    /// 打开表 + `id` 列索引（`id` 沿用老文件名 `表名.idx`）
     pub fn open_indexed<P: AsRef<Path>, Q: AsRef<Path>>(
         path: P,
         index_path: Q,
@@ -152,19 +181,155 @@ impl HeapTable {
         btree_order: usize,
         pool_size: usize,
     ) -> io::Result<Self> {
-        Ok(HeapTable {
+        let mut t = HeapTable {
             file: PageFile::with_options(path, page_size, pool_size)?,
-            index: Some(DiskBTree::open(index_path, page_size, btree_order, pool_size)?),
-        })
+            indexes: Vec::new(),
+        };
+        t.attach_index("id", index_path, page_size, btree_order, pool_size)?;
+        Ok(t)
     }
 
-    /// 插入一行，不更新索引
-    pub fn insert(&mut self, row: &Row) -> io::Result<()> {
-        self.insert_returning_position(row).map(|_| ())
+    /// 打开表 + 任意组索引：(列名, 索引文件路径)
+    pub fn open_with_indexes<P: AsRef<Path>>(
+        path: P,
+        indexes: &[(String, std::path::PathBuf)],
+        page_size: usize,
+        btree_order: usize,
+        pool_size: usize,
+    ) -> io::Result<Self> {
+        let mut t = HeapTable {
+            file: PageFile::with_options(path, page_size, pool_size)?,
+            indexes: Vec::new(),
+        };
+        for (col, p) in indexes {
+            t.attach_index(col, p, page_size, btree_order, pool_size)?;
+        }
+        Ok(t)
     }
 
-    /// 插入并返回行位置
-    pub fn insert_returning_position(&mut self, row: &Row) -> io::Result<(PageId, u16)> {
+    /// 挂上一棵索引树（只打开/创建文件，不填数据）
+    fn attach_index<P: AsRef<Path>>(
+        &mut self,
+        column: &str,
+        path: P,
+        page_size: usize,
+        btree_order: usize,
+        pool_size: usize,
+    ) -> io::Result<()> {
+        let tree = DiskBTree::open(path, page_size, btree_order, pool_size)?;
+        self.indexes.push(NamedIndex {
+            column: column.to_string(),
+            tree,
+        });
+        Ok(())
+    }
+
+    /// 这一列有索引吗
+    pub fn has_index(&self, column: &str) -> bool {
+        self.indexes.iter().any(|i| i.column == column)
+    }
+
+    /// 建索引：先扫一遍查重，再建树填数据。
+    ///
+    /// **要求这一列的整数取值唯一**，不唯一就报错、一个索引文件都不留。
+    /// 这样"走索引取一行"和"全表扫一遍再筛"必然得到同一个答案；
+    /// 非唯一索引要等 B+ 树支持重复键（记在项目计划的 P2-② 里）。
+    pub fn build_index(
+        &mut self,
+        column: &str,
+        path: &Path,
+        page_size: usize,
+        btree_order: usize,
+        pool_size: usize,
+    ) -> io::Result<()> {
+        if self.has_index(column) {
+            return Err(io::Error::new(
+                io::ErrorKind::AlreadyExists,
+                format!("index on column \"{}\" already exists", column),
+            ));
+        }
+
+        // 第 1 步：扫全表，查重并攒下 (键, 位置)
+        let mut seen: HashSet<i64> = HashSet::new();
+        let mut entries: Vec<(i64, u64)> = Vec::new();
+        for (page_id, slot, r) in self.scan_with_positions()? {
+            let Some(k) = row_int(&r, column) else {
+                continue; // 这一列不是整数的行不进索引
+            };
+            if !seen.insert(k) {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidInput,
+                    format!(
+                        "column \"{}\" has duplicate value {} — 暂不支持非唯一索引",
+                        column, k
+                    ),
+                ));
+            }
+            entries.push((k, pack_position(page_id, slot)));
+        }
+
+        // 第 2 步：建树
+        self.attach_index(column, path, page_size, btree_order, pool_size)?;
+        let i = self.indexes.len() - 1;
+        self.indexes[i].tree.clear()?; // 清掉可能残留的旧文件内容
+        for (k, pos) in entries {
+            self.indexes[i].tree.insert(k, pos)?;
+        }
+        Ok(())
+    }
+
+    /// 摘掉一个索引（不删文件）
+    pub fn detach_index(&mut self, column: &str) -> bool {
+        match self.indexes.iter().position(|i| i.column == column) {
+            None => false,
+            Some(i) => {
+                self.indexes.remove(i);
+                true
+            }
+        }
+    }
+
+    /// 插入一行：进堆表，并按每个索引的列顺手写索引
+    pub fn insert_row(&mut self, row: &Row) -> io::Result<()> {
+        self.insert_row_returning_position(row).map(|_| ())
+    }
+
+    /// 插入一行并返回它的位置
+    pub fn insert_row_returning_position(&mut self, row: &Row) -> io::Result<(PageId, u16)> {
+        // 先把键算好并查重：**查重必须在落堆之前**，
+        // 否则某个索引写不进去时行已经进堆了，表和索引就不一致了。
+        let mut pending: Vec<(usize, i64)> = Vec::new();
+        for (i, idx) in self.indexes.iter_mut().enumerate() {
+            let Some(k) = row_int(row, &idx.column) else {
+                continue; // 这一列不是整数（或没这一列）：这行不进这个索引
+            };
+            if idx.tree.get(k)?.is_some() {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidInput,
+                    format!("duplicate key {} on indexed column \"{}\"", k, idx.column),
+                ));
+            }
+            pending.push((i, k));
+        }
+
+        let (page_id, slot) = self.insert_returning_position(row)?;
+        let pos = pack_position(page_id, slot);
+        for (i, k) in pending {
+            self.indexes[i].tree.insert(k, pos)?;
+        }
+        Ok((page_id, slot))
+    }
+
+    /// 批量插入（只写缓冲池，落盘由调用方 `flush` 一次完成）
+    pub fn insert_rows(&mut self, rows: &[Row]) -> io::Result<()> {
+        for r in rows {
+            self.insert_row(r)?;
+        }
+        Ok(())
+    }
+
+    /// 插入一行，只进堆表、不碰索引
+    fn insert_returning_position(&mut self, row: &Row) -> io::Result<(PageId, u16)> {
         let tuple = serde_json::to_vec(row).map_err(io::Error::other)?;
 
         if HEADER_SIZE + SLOT_SIZE + tuple.len() > self.file.page_size() {
@@ -193,23 +358,18 @@ impl HeapTable {
         Ok((page_id, slot))
     }
 
-    /// 插入一行并写索引
-    pub fn insert_keyed(&mut self, row: &Row, key: i64) -> io::Result<()> {
-        let (page_id, slot) = self.insert_returning_position(row)?;
-        if let Some(idx) = &mut self.index {
-            idx.insert(key, pack_position(page_id, slot))?;
-        }
-        Ok(())
+    /// 按 `id` 取一行
+    pub fn get_by_key(&mut self, key: i64) -> io::Result<Option<Row>> {
+        self.get_by_column_key("id", key)
     }
 
-    /// 按索引键取一行
-    pub fn get_by_key(&mut self, key: i64) -> io::Result<Option<Row>> {
-        let pos = match &mut self.index {
-            None => return Ok(None),
-            Some(idx) => match idx.get(key)? {
-                None => return Ok(None),
-                Some(p) => p,
-            },
+    /// 按某一列的索引取一行
+    pub fn get_by_column_key(&mut self, column: &str, key: i64) -> io::Result<Option<Row>> {
+        let Some(idx) = self.indexes.iter_mut().find(|i| i.column == column) else {
+            return Ok(None);
+        };
+        let Some(pos) = idx.tree.get(key)? else {
+            return Ok(None);
         };
         let (page_id, slot) = unpack_position(pos);
         self.read_at(page_id, slot)
@@ -227,47 +387,115 @@ impl HeapTable {
         }
     }
 
-    /// 全表扫描
-    pub fn scan(&mut self) -> io::Result<Vec<Row>> {
-        let n = self.file.num_pages()?;
-        let mut rows = Vec::new();
-        for i in 0..n {
-            let tuples = self.file.with_page(i, |p| {
-                let mut out = Vec::new();
-                for slot in HeapPage::iter_slots(p) {
-                    if let Some(bytes) = HeapPage::get_tuple(p, slot) {
-                        out.push(bytes);
-                    }
+    /// 按位置删一行：槽位长度置 0，从所有索引里摘掉，并把这一行交出来
+    pub fn delete_at(&mut self, page_id: PageId, slot: u16) -> io::Result<Option<Row>> {
+        // 先读出来：删完就不知道这行各列是什么，索引也就无从摘起
+        let row = self.read_at(page_id, slot)?;
+        let existed = self.file.update_page(page_id, |p| HeapPage::delete_tuple(p, slot))?;
+        if !existed {
+            return Ok(None);
+        }
+        if let Some(r) = &row {
+            for idx in &mut self.indexes {
+                if let Some(k) = row_int(r, &idx.column) {
+                    idx.tree.delete(k)?;
                 }
-                out
-            })?;
-            for bytes in tuples {
-                let row: Row = serde_json::from_slice(&bytes).map_err(io::Error::other)?;
-                rows.push(row);
             }
         }
-        Ok(rows)
+        Ok(row)
     }
 
-    /// 整表改写并重建索引
+    /// 按 `id` 批量删行，返回被删掉的那些行（数据字典要用它更新统计）。
+    ///
+    /// 定位靠 `id` 索引；没有这个索引就退化成一趟全表扫描（仍然只删目标行）。
+    pub fn delete_by_keys(&mut self, keys: &[i64]) -> io::Result<Vec<Row>> {
+        let mut positions: Vec<(PageId, u16)> = Vec::new();
+
+        if self.has_index("id") {
+            for &k in keys {
+                let pos = self
+                    .indexes
+                    .iter_mut()
+                    .find(|i| i.column == "id")
+                    .and_then(|i| i.tree.get(k).ok().flatten());
+                if let Some(p) = pos {
+                    positions.push(unpack_position(p));
+                }
+            }
+        } else {
+            let mut wanted: Vec<i64> = keys.to_vec();
+            wanted.sort_unstable();
+            for (page_id, slot, r) in self.scan_with_positions()? {
+                if let Some(k) = row_int(&r, "id") {
+                    if wanted.binary_search(&k).is_ok() {
+                        positions.push((page_id, slot));
+                    }
+                }
+            }
+        }
+
+        let mut deleted: Vec<Row> = Vec::new();
+        for (page_id, slot) in positions {
+            if let Some(r) = self.delete_at(page_id, slot)? {
+                deleted.push(r);
+            }
+        }
+        Ok(deleted)
+    }
+
+    /// 全表扫描
+    pub fn scan(&mut self) -> io::Result<Vec<Row>> {
+        Ok(self
+            .scan_with_positions()?
+            .into_iter()
+            .map(|(_, _, r)| r)
+            .collect())
+    }
+
+    /// 全表扫描，同时给出每行的位置（建索引用）
+    pub fn scan_with_positions(&mut self) -> io::Result<Vec<(PageId, u16, Row)>> {
+        let n = self.file.num_pages()?;
+        let mut out = Vec::new();
+        for i in 0..n {
+            let tuples = self.file.with_page(i, |p| {
+                let mut v = Vec::new();
+                for slot in HeapPage::iter_slots(p) {
+                    if let Some(bytes) = HeapPage::get_tuple(p, slot) {
+                        v.push((slot, bytes));
+                    }
+                }
+                v
+            })?;
+            for (slot, bytes) in tuples {
+                let row: Row = serde_json::from_slice(&bytes).map_err(io::Error::other)?;
+                out.push((i, slot, row));
+            }
+        }
+        Ok(out)
+    }
+
+    /// 整表改写并重建所有索引
     pub fn replace_all(&mut self, rows: &[Row]) -> io::Result<()> {
         self.file.truncate()?;
-        if let Some(idx) = &mut self.index {
-            idx.clear()?;
+        for idx in &mut self.indexes {
+            idx.tree.clear()?;
         }
         for r in rows {
             let (page_id, slot) = self.insert_returning_position(r)?;
-            if let (Some(idx), Some(key)) = (&mut self.index, row_key(r)) {
-                idx.insert(key, pack_position(page_id, slot))?;
+            let pos = pack_position(page_id, slot);
+            for idx in &mut self.indexes {
+                if let Some(k) = row_int(r, &idx.column) {
+                    idx.tree.insert(k, pos)?;
+                }
             }
         }
         Ok(())
     }
 
-    /// 把脏页写回
+    /// 把脏页写回（索引也一起）
     pub fn flush(&mut self) -> io::Result<()> {
-        if let Some(idx) = &mut self.index {
-            idx.flush()?;
+        for idx in &mut self.indexes {
+            idx.tree.flush()?;
         }
         self.file.flush()
     }
@@ -282,11 +510,3 @@ impl HeapTable {
     /// 命中率
     pub fn hit_rate(&self) -> f64 { self.file.hit_rate() }
 }
-
-/// 取行里的 id 作索引键
-fn row_key(row: &Row) -> Option<i64> {
-    match row.get("id") {
-        Some(serde_json::Value::Number(n)) => n.as_i64(),
-        _ => None,
-    }
-}