optimize query execution and caching

pkg/executor/executor.go

@@ -1,6 +1,7 @@
 package executor
 
 import (
+	"container/heap"
 	"errors"
 	"fmt"
 	"math/rand"
@@ -244,6 +245,44 @@ func isCatalogMiss(err error) bool {
 		analysisErr.Type == analyzer.ErrColumnNotFound
 }
 
+// isCountStarSingleTable reports whether the statement is the safe COUNT(*)
+// shape eligible for the metadata fast path: a single-table SELECT with exactly
+// one COUNT(*) column and no filters, grouping, DISTINCT, JOIN, subquery, or
+// LIMIT/OFFSET. Anything else returns false so unsupported shapes use the
+// normal scan path.
+func isCountStarSingleTable(stmt *parser.SelectStmt) bool {
+	if stmt.Compound != nil || stmt.Distinct {
+		return false
+	}
+	if stmt.Where != nil || stmt.Having != nil {
+		return false
+	}
+	if len(stmt.GroupBy) > 0 || len(stmt.OrderBy) > 0 {
+		return false
+	}
+	if stmt.Limit != nil || stmt.Offset != nil {
+		return false
+	}
+	if len(stmt.From) != 1 {
+		return false
+	}
+	ref := stmt.From[0]
+	if ref.Subquery != nil || ref.Join != nil {
+		return false
+	}
+	if len(stmt.Columns) != 1 || stmt.Columns[0].Star {
+		return false
+	}
+	fn, ok := stmt.Columns[0].Expr.(*parser.FunctionCall)
+	if !ok {
+		return false
+	}
+	if !strings.EqualFold(fn.Name, "count") || !fn.Star || len(fn.Args) > 0 {
+		return false
+	}
+	return true
+}
+
 // executeSelect executes a SELECT statement (or compound SELECT).
 func (e *Executor) executeSelect(stmt *parser.SelectStmt) (*Result, error) {
 	if stmt.Compound != nil {
@@ -280,6 +319,24 @@ func (e *Executor) executeSelect(stmt *parser.SelectStmt) (*Result, error) {
 		return nil, err
 	}
 
+	// COUNT(*) fast path: exact metadata-based count for the safe single-table
+	// shape with no filters/grouping/distinct/join. Any unsupported shape falls
+	// through to the normal scan path.
+	if isCountStarSingleTable(stmt) {
+		count, err := e.table.CountFast(tableName)
+		if err != nil {
+			return nil, err
+		}
+		result := NewResult("SELECT")
+		if stmt.Columns[0].Alias != "" {
+			result.AddColumn(stmt.Columns[0].Alias)
+		} else {
+			result.AddColumn("column1")
+		}
+		result.AddRow(int64(count))
+		return result, nil
+	}
+
 	// Multi-table FROM (comma-separated implicit cross join): collect and cross join all tables,
 	// then apply WHERE after. Don't push WHERE down here — conditions reference multiple tables.
 	isMultiTable := len(stmt.From) > 1 && stmt.From[0].Join == nil
@@ -999,26 +1056,8 @@ func (e *Executor) executeSelect(stmt *parser.SelectStmt) (*Result, error) {
 		return e.executeAggregateSelect(stmt, rows, schema)
 	}
 
-	// Apply ORDER BY
-	if len(stmt.OrderBy) > 0 {
-		e.sortRows(rows, resolveOrderByPositions(stmt.OrderBy, stmt.Columns))
-	}
-
-	// Apply LIMIT/OFFSET
-	if stmt.Offset != nil {
-		offset := e.evalIntExpr(stmt.Offset)
-		if offset < len(rows) {
-			rows = rows[offset:]
-		} else {
-			rows = nil
-		}
-	}
-	if stmt.Limit != nil {
-		limit := e.evalIntExpr(stmt.Limit)
-		if limit < len(rows) {
-			rows = rows[:limit]
-		}
-	}
+	// Apply ORDER BY, LIMIT, and OFFSET.
+	rows = e.orderAndLimitRows(rows, stmt.OrderBy, stmt.Limit, stmt.Offset, stmt.Columns)
 
 	// Build result
 	result := NewResult("SELECT")
@@ -1318,26 +1357,8 @@ func (e *Executor) executeSelectFromSubquery(stmt *parser.SelectStmt) (*Result,
 		return e.executeAggregateSelect(stmt, derivedRows, tempSchema)
 	}
 
-	// Apply ORDER BY
-	if len(stmt.OrderBy) > 0 {
-		e.sortRows(derivedRows, resolveOrderByPositions(stmt.OrderBy, stmt.Columns))
-	}
-
-	// Apply LIMIT/OFFSET
-	if stmt.Offset != nil {
-		offset := e.evalIntExpr(stmt.Offset)
-		if offset < len(derivedRows) {
-			derivedRows = derivedRows[offset:]
-		} else {
-			derivedRows = nil
-		}
-	}
-	if stmt.Limit != nil {
-		limit := e.evalIntExpr(stmt.Limit)
-		if limit < len(derivedRows) {
-			derivedRows = derivedRows[:limit]
-		}
-	}
+	// Apply ORDER BY, LIMIT, and OFFSET.
+	derivedRows = e.orderAndLimitRows(derivedRows, stmt.OrderBy, stmt.Limit, stmt.Offset, stmt.Columns)
 
 	// Build result
 	result := NewResult("SELECT")
@@ -1701,25 +1722,7 @@ func (e *Executor) finalizeGroupResult(stmt *parser.SelectStmt, result *Result,
 	if stmt.Distinct {
 		result.Rows = e.applyDistinct(result.Rows)
 	}
-	if len(stmt.OrderBy) > 0 {
-		e.sortResultRows(result, stmt.OrderBy, expandedColumns, columnNames)
-	}
-	if stmt.Offset != nil {
-		offset := e.evalIntExpr(stmt.Offset)
-		if offset < len(result.Rows) {
-			result.Rows = result.Rows[offset:]
-		} else {
-			result.Rows = nil
-		}
-		result.RowCount = len(result.Rows)
-	}
-	if stmt.Limit != nil {
-		limit := e.evalIntExpr(stmt.Limit)
-		if limit < len(result.Rows) {
-			result.Rows = result.Rows[:limit]
-		}
-		result.RowCount = len(result.Rows)
-	}
+	e.orderAndLimitResultRows(result, stmt.OrderBy, stmt.Limit, stmt.Offset, expandedColumns, columnNames)
 	return result, nil
 }
 
@@ -4970,64 +4973,254 @@ func resolveOrderByPositions(orderBy []parser.OrderByItem, selectCols []parser.S
 	return result
 }
 
+// orderByLess reports whether key slice a sorts before b under orderBy.
+// Keys are precomputed per-row ORDER BY expression values, one per item.
+func orderByLess(a, b []interface{}, orderBy []parser.OrderByItem) bool {
+	for i, item := range orderBy {
+		cmp := compare(a[i], b[i])
+		if cmp != 0 {
+			if item.Desc {
+				return cmp > 0
+			}
+			return cmp < 0
+		}
+	}
+	return false
+}
+
+// sortKeyOrder returns a permutation of [0..len(keys)) that sorts the keys
+// ascending per orderBy.
+func sortKeyOrder(keys [][]interface{}, orderBy []parser.OrderByItem) []int {
+	order := make([]int, len(keys))
+	for i := range order {
+		order[i] = i
+	}
+	sort.Slice(order, func(a, b int) bool {
+		return orderByLess(keys[order[a]], keys[order[b]], orderBy)
+	})
+	return order
+}
+
+// topNHeap is a bounded max-heap that keeps the k smallest elements (per
+// orderByLess) seen so far.
+type topNHeap struct {
+	keys    [][]interface{}
+	idx     []int
+	orderBy []parser.OrderByItem
+}
+
+func (h *topNHeap) Len() int { return len(h.idx) }
+
+func (h *topNHeap) Less(i, j int) bool {
+	return orderByLess(h.keys[h.idx[j]], h.keys[h.idx[i]], h.orderBy)
+}
+
+func (h *topNHeap) Swap(i, j int) { h.idx[i], h.idx[j] = h.idx[j], h.idx[i] }
+
+func (h *topNHeap) Push(x interface{}) { h.idx = append(h.idx, x.(int)) }
+
+func (h *topNHeap) Pop() interface{} {
+	n := len(h.idx)
+	x := h.idx[n-1]
+	h.idx = h.idx[:n-1]
+	return x
+}
+
+// topNKeyOrder returns the indices of the k smallest keys (per orderByLess) in
+// ascending order, without fully sorting all n elements. If k >= n it falls
+// back to a full sort.
+func topNKeyOrder(keys [][]interface{}, orderBy []parser.OrderByItem, k int) []int {
+	if k <= 0 {
+		return nil
+	}
+	if k >= len(keys) {
+		return sortKeyOrder(keys, orderBy)
+	}
+	h := &topNHeap{keys: keys, orderBy: orderBy, idx: make([]int, 0, k)}
+	for i := range keys {
+		if h.Len() < k {
+			heap.Push(h, i)
+		} else if orderByLess(keys[i], keys[h.idx[0]], orderBy) {
+			h.idx[0] = i
+			heap.Fix(h, 0)
+		}
+	}
+	selected := append([]int(nil), h.idx...)
+	sort.Slice(selected, func(a, b int) bool {
+		return orderByLess(keys[selected[a]], keys[selected[b]], orderBy)
+	})
+	return selected
+}
+
+func reorderRows(rows []storage.Row, order []int) {
+	tmp := make([]storage.Row, len(rows))
+	for i, idx := range order {
+		tmp[i] = rows[idx]
+	}
+	copy(rows, tmp)
+}
+
+// sortRowKeys precomputes the ORDER BY expression value for each row so each
+// expression is evaluated once per row instead of O(n log n) times.
+func (e *Executor) sortRowKeys(rows []storage.Row, orderBy []parser.OrderByItem) [][]interface{} {
+	keys := make([][]interface{}, len(rows))
+	for i, row := range rows {
+		ks := make([]interface{}, len(orderBy))
+		for j, item := range orderBy {
+			ks[j], _ = e.evalExpr(item.Expr, row)
+		}
+		keys[i] = ks
+	}
+	return keys
+}
+
 func (e *Executor) sortRows(rows []storage.Row, orderBy []parser.OrderByItem) {
-	sort.Slice(rows, func(i, j int) bool {
-		for _, item := range orderBy {
-			vi, _ := e.evalExpr(item.Expr, rows[i])
-			vj, _ := e.evalExpr(item.Expr, rows[j])
-			cmp := compare(vi, vj)
-			if cmp != 0 {
-				if item.Desc {
-					return cmp > 0
+	order := sortKeyOrder(e.sortRowKeys(rows, orderBy), orderBy)
+	reorderRows(rows, order)
+}
+
+// topNRows sorts only enough to keep the k smallest rows (per orderBy).
+func (e *Executor) topNRows(rows []storage.Row, orderBy []parser.OrderByItem, k int) []storage.Row {
+	order := topNKeyOrder(e.sortRowKeys(rows, orderBy), orderBy, k)
+	out := make([]storage.Row, len(order))
+	for i, idx := range order {
+		out[i] = rows[idx]
+	}
+	return out
+}
+
+// orderAndLimitRows applies ORDER BY (with a bounded top-N selection when a
+// LIMIT is present), then OFFSET and LIMIT, preserving SQL semantics.
+func (e *Executor) orderAndLimitRows(rows []storage.Row, orderBy []parser.OrderByItem, limitExpr, offsetExpr parser.Expr, selectCols []parser.SelectColumn) []storage.Row {
+	orderBy = resolveOrderByPositions(orderBy, selectCols)
+
+	var offset, limit int
+	hasOffset := offsetExpr != nil
+	hasLimit := limitExpr != nil
+	if hasOffset {
+		offset = e.evalIntExpr(offsetExpr)
+	}
+	if hasLimit {
+		limit = e.evalIntExpr(limitExpr)
+	}
+
+	if len(orderBy) > 0 {
+		if hasLimit && limit >= 0 {
+			k := offset + limit
+			if k >= 0 && k < len(rows) {
+				rows = e.topNRows(rows, orderBy, k)
+			} else {
+				e.sortRows(rows, orderBy)
+			}
+		} else {
+			e.sortRows(rows, orderBy)
+		}
+	}
+
+	if hasOffset {
+		if offset < len(rows) {
+			rows = rows[offset:]
+		} else {
+			rows = nil
+		}
+	}
+	if hasLimit {
+		if limit < len(rows) {
+			rows = rows[:limit]
+		}
+	}
+	return rows
+}
+
+// resultRowKey evaluates a single ORDER BY item against a result row, honoring
+// select-column aliases exactly like the previous sortResultRows implementation.
+func (e *Executor) resultRowKey(result *Result, rowIdx int, item parser.OrderByItem, columnNames []string) interface{} {
+	if ref, ok := item.Expr.(*parser.ColumnRef); ok && ref.Table == "" {
+		for idx, name := range columnNames {
+			if strings.EqualFold(name, ref.Column) {
+				if idx < len(result.Rows[rowIdx]) {
+					return result.Rows[rowIdx][idx]
 				}
-				return cmp < 0
 			}
 		}
-		return false
-	})
+	}
+	row := e.resultRowToStorageRow(result, rowIdx)
+	v, _ := e.evalExpr(item.Expr, row)
+	return v
+}
+
+// resultRowKeys precomputes the ORDER BY expression value for each result row.
+func (e *Executor) resultRowKeys(result *Result, orderBy []parser.OrderByItem, columnNames []string) [][]interface{} {
+	keys := make([][]interface{}, len(result.Rows))
+	for i := range result.Rows {
+		ks := make([]interface{}, len(orderBy))
+		for j, item := range orderBy {
+			ks[j] = e.resultRowKey(result, i, item, columnNames)
+		}
+		keys[i] = ks
+	}
+	return keys
 }
 
 // sortResultRows sorts Result.Rows based on ORDER BY clauses.
 // It handles column aliases by matching them against the select columns.
 func (e *Executor) sortResultRows(result *Result, orderBy []parser.OrderByItem, selectColumns []parser.SelectColumn, columnNames []string) {
 	orderBy = resolveOrderByPositions(orderBy, selectColumns)
-	sort.Slice(result.Rows, func(i, j int) bool {
-		for _, item := range orderBy {
-			var vi, vj interface{}
-			var rowI, rowJ storage.Row
-
-			// Check if ORDER BY references a column alias
-			if ref, ok := item.Expr.(*parser.ColumnRef); ok && ref.Table == "" {
-				// Look for matching alias in select columns
-				for idx, name := range columnNames {
-					if strings.EqualFold(name, ref.Column) {
-						if idx < len(result.Rows[i]) {
-							vi = result.Rows[i][idx]
-							vj = result.Rows[j][idx]
-							goto compare
-						}
-					}
-				}
-			}
-
-			// If not found as alias, try to evaluate the expression
-			// Create temporary rows from result rows for evaluation
-			rowI = e.resultRowToStorageRow(result, i)
-			rowJ = e.resultRowToStorageRow(result, j)
-			vi, _ = e.evalExpr(item.Expr, rowI)
-			vj, _ = e.evalExpr(item.Expr, rowJ)
+	order := sortKeyOrder(e.resultRowKeys(result, orderBy, columnNames), orderBy)
+	rows := make([][]interface{}, len(order))
+	for i, idx := range order {
+		rows[i] = result.Rows[idx]
+	}
+	result.Rows = rows
+}
 
-		compare:
-			cmp := compare(vi, vj)
-			if cmp != 0 {
-				if item.Desc {
-					return cmp > 0
+// orderAndLimitResultRows applies ORDER BY (with bounded top-N selection when a
+// LIMIT is present), then OFFSET and LIMIT, to a Result's rows.
+func (e *Executor) orderAndLimitResultRows(result *Result, orderBy []parser.OrderByItem, limitExpr, offsetExpr parser.Expr, selectCols []parser.SelectColumn, columnNames []string) {
+	orderBy = resolveOrderByPositions(orderBy, selectCols)
+
+	var offset, limit int
+	hasOffset := offsetExpr != nil
+	hasLimit := limitExpr != nil
+	if hasOffset {
+		offset = e.evalIntExpr(offsetExpr)
+	}
+	if hasLimit {
+		limit = e.evalIntExpr(limitExpr)
+	}
+
+	if len(orderBy) > 0 {
+		if hasLimit && limit >= 0 {
+			k := offset + limit
+			if k >= 0 && k < len(result.Rows) {
+				order := topNKeyOrder(e.resultRowKeys(result, orderBy, columnNames), orderBy, k)
+				rows := make([][]interface{}, len(order))
+				for i, idx := range order {
+					rows[i] = result.Rows[idx]
 				}
-				return cmp < 0
+				result.Rows = rows
+			} else {
+				e.sortResultRows(result, orderBy, nil, columnNames)
 			}
+		} else {
+			e.sortResultRows(result, orderBy, nil, columnNames)
 		}
-		return false
-	})
+	}
+
+	if hasOffset {
+		if offset < len(result.Rows) {
+			result.Rows = result.Rows[offset:]
+		} else {
+			result.Rows = nil
+		}
+		result.RowCount = len(result.Rows)
+	}
+	if hasLimit {
+		if limit < len(result.Rows) {
+			result.Rows = result.Rows[:limit]
+		}
+		result.RowCount = len(result.Rows)
+	}
 }
 
 // resultRowToStorageRow converts a Result row back to storage.Row for expression evaluation.

pkg/executor/optimizations_test.go

@@ -0,0 +1,329 @@
+package executor
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/danfragoso/pizzasql-next/pkg/storage"
+)
+
+// newOptExec creates an executor backed by PizzaKV, skipping the test when it
+// is unavailable (mirroring the convention in executor_test.go).
+func newOptExec(t *testing.T, db string) *Executor {
+	t.Helper()
+	pool, err := storage.NewKVPool("localhost:8085", 5, 5*time.Second)
+	if err != nil {
+		t.Skipf("PizzaKV not available: %v", err)
+	}
+	t.Cleanup(func() { pool.Close() })
+	schema := storage.NewSchemaManager(pool, db)
+	table := storage.NewTableManager(pool, schema, db)
+	return New(schema, table)
+}
+
+// intColumn extracts a column of int64 values from a result.
+func intColumn(t *testing.T, r *Result, col int) []int64 {
+	t.Helper()
+	out := make([]int64, 0, len(r.Rows))
+	for _, row := range r.Rows {
+		if col >= len(row) {
+			t.Fatalf("row too short: %v", row)
+		}
+		out = append(out, row[col].(int64))
+	}
+	return out
+}
+
+// TestTopNOrderByLimitMatchesFullSort verifies bounded top-N execution returns
+// exactly the same rows (order included) as a full sort followed by LIMIT/OFFSET.
+func TestTopNOrderByLimitMatchesFullSort(t *testing.T) {
+	exec := newOptExec(t, "test_topn_db")
+	execSQL(exec, "DROP TABLE IF EXISTS nums")
+	if _, err := execSQL(exec, "CREATE TABLE nums (id INTEGER PRIMARY KEY, v INTEGER)"); err != nil {
+		t.Fatalf("create: %v", err)
+	}
+	defer execSQL(exec, "DROP TABLE IF EXISTS nums")
+
+	const n = 500
+	for i := 0; i < n; i++ {
+		// Deterministic permutation of 0..n-1.
+		v := (i*137 + 41) % n
+		if _, err := execSQL(exec, fmt.Sprintf("INSERT INTO nums VALUES (%d, %d)", i+1, v)); err != nil {
+			t.Fatalf("insert: %v", err)
+		}
+	}
+
+	for _, desc := range []bool{false, true} {
+		for _, offset := range []int{0, 3, 47, n - 1} {
+			for _, limit := range []int{1, 2, 17, 100, n + 5} {
+				dir := "ASC"
+				if desc {
+					dir = "DESC"
+				}
+				fullRes, err := execSQL(exec, fmt.Sprintf("SELECT v FROM nums ORDER BY v %s", dir))
+				if err != nil {
+					t.Fatalf("full: %v", err)
+				}
+				full := intColumn(t, fullRes, 0)
+
+				q := fmt.Sprintf("SELECT v FROM nums ORDER BY v %s LIMIT %d OFFSET %d", dir, limit, offset)
+				limRes, err := execSQL(exec, q)
+				if err != nil {
+					t.Fatalf("%s: %v", q, err)
+				}
+				got := intColumn(t, limRes, 0)
+
+				want := sliceRange(full, offset, limit)
+				if !equalInt64s(got, want) {
+					t.Fatalf("%s: got %v want %v", q, got, want)
+				}
+			}
+		}
+	}
+}
+
+// TestTopNOrderByLimitTies verifies LIMIT/OFFSET with tied ORDER BY keys returns
+// the correct multiset of values even though tie ordering is unspecified.
+func TestTopNOrderByLimitTies(t *testing.T) {
+	exec := newOptExec(t, "test_topn_ties_db")
+	execSQL(exec, "DROP TABLE IF EXISTS ties")
+	if _, err := execSQL(exec, "CREATE TABLE ties (id INTEGER PRIMARY KEY, v INTEGER)"); err != nil {
+		t.Fatalf("create: %v", err)
+	}
+	defer execSQL(exec, "DROP TABLE IF EXISTS ties")
+
+	// 4 rows with v=1, 2 rows with v=2, 1 row with v=3.
+	vals := []int{1, 1, 1, 1, 2, 2, 3}
+	for i, v := range vals {
+		if _, err := execSQL(exec, fmt.Sprintf("INSERT INTO ties VALUES (%d, %d)", i+1, v)); err != nil {
+			t.Fatalf("insert: %v", err)
+		}
+	}
+
+	// LIMIT 3: the three smallest, all v=1.
+	res, err := execSQL(exec, "SELECT v FROM ties ORDER BY v LIMIT 3")
+	if err != nil {
+		t.Fatalf("limit 3: %v", err)
+	}
+	got := intColumn(t, res, 0)
+	if len(got) != 3 || got[0] != 1 || got[1] != 1 || got[2] != 1 {
+		t.Fatalf("LIMIT 3 got %v, want [1 1 1]", got)
+	}
+
+	// OFFSET 3 LIMIT 3: skip three v=1 rows, then one v=1 + two v=2.
+	res, err = execSQL(exec, "SELECT v FROM ties ORDER BY v LIMIT 3 OFFSET 3")
+	if err != nil {
+		t.Fatalf("offset 3 limit 3: %v", err)
+	}
+	got = intColumn(t, res, 0)
+	if len(got) != 3 || got[0] != 1 || got[1] != 2 || got[2] != 2 {
+		t.Fatalf("OFFSET 3 LIMIT 3 got %v, want [1 2 2]", got)
+	}
+
+	// OFFSET beyond the ties boundary.
+	res, err = execSQL(exec, "SELECT v FROM ties ORDER BY v LIMIT 2 OFFSET 5")
+	if err != nil {
+		t.Fatalf("offset 5 limit 2: %v", err)
+	}
+	got = intColumn(t, res, 0)
+	if len(got) != 2 || got[0] != 2 || got[1] != 3 {
+		t.Fatalf("OFFSET 5 LIMIT 2 got %v, want [2 3]", got)
+	}
+}
+
+// TestTopNGroupByOrderLimit verifies the grouped/result-row top-N path.
+func TestTopNGroupByOrderLimit(t *testing.T) {
+	exec := newOptExec(t, "test_topn_group_db")
+	execSQL(exec, "DROP TABLE IF EXISTS sales")
+	if _, err := execSQL(exec, "CREATE TABLE sales (id INTEGER PRIMARY KEY, region TEXT, amount INTEGER)"); err != nil {
+		t.Fatalf("create: %v", err)
+	}
+	defer execSQL(exec, "DROP TABLE IF EXISTS sales")
+
+	for _, r := range []struct {
+		id     int
+		region string
+		amount int
+	}{
+		{1, "east", 10}, {2, "west", 20}, {3, "north", 30},
+		{4, "south", 40}, {5, "east", 50},
+	} {
+		if _, err := execSQL(exec, fmt.Sprintf("INSERT INTO sales VALUES (%d, '%s', %d)", r.id, r.region, r.amount)); err != nil {
+			t.Fatalf("insert: %v", err)
+		}
+	}
+
+	res, err := execSQL(exec, "SELECT region, COUNT(*) AS c FROM sales GROUP BY region ORDER BY region LIMIT 2")
+	if err != nil {
+		t.Fatalf("group topn: %v", err)
+	}
+	if len(res.Rows) != 2 {
+		t.Fatalf("expected 2 rows, got %d: %v", len(res.Rows), res.Rows)
+	}
+	if res.Rows[0][0] != "east" || res.Rows[1][0] != "north" {
+		t.Fatalf("unexpected top-2 regions: %v", res.Rows)
+	}
+	if res.Rows[0][1] != int64(2) {
+		t.Fatalf("unexpected east count: %v", res.Rows[0][1])
+	}
+}
+
+// TestCountFastPathLifecycleAndRollback verifies the exact COUNT(*) fast path
+// across writes and transaction rollback.
+func TestCountFastPathLifecycleAndRollback(t *testing.T) {
+	exec := newOptExec(t, "test_count_db")
+	execSQL(exec, "DROP TABLE IF EXISTS items")
+	if _, err := execSQL(exec, "CREATE TABLE items (id INTEGER PRIMARY KEY, v TEXT)"); err != nil {
+		t.Fatalf("create: %v", err)
+	}
+	defer execSQL(exec, "DROP TABLE IF EXISTS items")
+
+	countStar := func() int64 {
+		t.Helper()
+		res, err := execSQL(exec, "SELECT COUNT(*) FROM items")
+		if err != nil {
+			t.Fatalf("count: %v", err)
+		}
+		return res.Rows[0][0].(int64)
+	}
+
+	if got := countStar(); got != 0 {
+		t.Fatalf("initial count = %d, want 0", got)
+	}
+	for i := 1; i <= 5; i++ {
+		if _, err := execSQL(exec, fmt.Sprintf("INSERT INTO items VALUES (%d, 'x%d')", i, i)); err != nil {
+			t.Fatalf("insert: %v", err)
+		}
+	}
+	if got := countStar(); got != 5 {
+		t.Fatalf("after inserts = %d, want 5", got)
+	}
+
+	// UPDATE keeps the count exact.
+	if _, err := execSQL(exec, "UPDATE items SET v = 'y' WHERE id = 1"); err != nil {
+		t.Fatalf("update: %v", err)
+	}
+	if got := countStar(); got != 5 {
+		t.Fatalf("after update = %d, want 5", got)
+	}
+
+	// Bulk insert via INSERT ... SELECT.
+	if _, err := execSQL(exec, "INSERT INTO items (id, v) SELECT id + 100, v FROM items"); err != nil {
+		t.Fatalf("insert-select: %v", err)
+	}
+	if got := countStar(); got != 10 {
+		t.Fatalf("after insert-select = %d, want 10", got)
+	}
+
+	// DELETE decrements.
+	if _, err := execSQL(exec, "DELETE FROM items WHERE id > 100"); err != nil {
+		t.Fatalf("delete: %v", err)
+	}
+	if got := countStar(); got != 5 {
+		t.Fatalf("after delete = %d, want 5", got)
+	}
+
+	// Transaction rollback restores the exact count.
+	execSQL(exec, "BEGIN")
+	if _, err := execSQL(exec, "INSERT INTO items VALUES (999, 'tmp')"); err != nil {
+		t.Fatalf("tx insert: %v", err)
+	}
+	if got := countStar(); got != 6 {
+		t.Fatalf("inside tx = %d, want 6", got)
+	}
+	if _, err := execSQL(exec, "ROLLBACK"); err != nil {
+		t.Fatalf("rollback: %v", err)
+	}
+	if got := countStar(); got != 5 {
+		t.Fatalf("after rollback = %d, want 5", got)
+	}
+}
+
+// TestCountFastPathRestart verifies a second executor (fresh process state)
+// derives the same exact count from durable rows.
+func TestCountFastPathRestart(t *testing.T) {
+	exec := newOptExec(t, "test_count_restart_db")
+	execSQL(exec, "DROP TABLE IF EXISTS r")
+	if _, err := execSQL(exec, "CREATE TABLE r (id INTEGER PRIMARY KEY)"); err != nil {
+		t.Fatalf("create: %v", err)
+	}
+	defer execSQL(exec, "DROP TABLE IF EXISTS r")
+	for i := 1; i <= 7; i++ {
+		if _, err := execSQL(exec, fmt.Sprintf("INSERT INTO r VALUES (%d)", i)); err != nil {
+			t.Fatalf("insert: %v", err)
+		}
+	}
+
+	// A brand-new executor over the same KV.
+	exec2 := newOptExec(t, "test_count_restart_db")
+	res, err := execSQL(exec2, "SELECT COUNT(*) FROM r")
+	if err != nil {
+		t.Fatalf("count after restart: %v", err)
+	}
+	if got := res.Rows[0][0].(int64); got != 7 {
+		t.Fatalf("count after restart = %d, want 7", got)
+	}
+}
+
+// TestCountFastPathUnsupportedShapesStillCorrect verifies shapes outside the
+// fast path fall through to the normal scan and produce correct results.
+func TestCountFastPathUnsupportedShapesStillCorrect(t *testing.T) {
+	exec := newOptExec(t, "test_count_unsupported_db")
+	execSQL(exec, "DROP TABLE IF EXISTS t2")
+	execSQL(exec, "DROP TABLE IF EXISTS t1")
+	if _, err := execSQL(exec, "CREATE TABLE t1 (id INTEGER PRIMARY KEY, x INTEGER)"); err != nil {
+		t.Fatalf("create t1: %v", err)
+	}
+	if _, err := execSQL(exec, "CREATE TABLE t2 (id INTEGER PRIMARY KEY, y INTEGER)"); err != nil {
+		t.Fatalf("create t2: %v", err)
+	}
+	defer execSQL(exec, "DROP TABLE IF EXISTS t2")
+	defer execSQL(exec, "DROP TABLE IF EXISTS t1")
+
+	for i := 1; i <= 4; i++ {
+		execSQL(exec, fmt.Sprintf("INSERT INTO t1 VALUES (%d, %d)", i, i))
+		execSQL(exec, fmt.Sprintf("INSERT INTO t2 VALUES (%d, %d)", i, i))
+	}
+
+	cases := []struct {
+		q    string
+		want int64
+	}{
+		{"SELECT COUNT(*) FROM t1 WHERE x > 2", 2},
+		{"SELECT COUNT(DISTINCT x) FROM t1", 4},
+		{"SELECT COUNT(*) FROM t1 t1a, t2 t2b", 16},
+	}
+	for _, c := range cases {
+		res, err := execSQL(exec, c.q)
+		if err != nil {
+			t.Fatalf("%s: %v", c.q, err)
+		}
+		if got := res.Rows[0][0].(int64); got != c.want {
+			t.Fatalf("%s = %d, want %d", c.q, got, c.want)
+		}
+	}
+}
+
+func sliceRange(v []int64, offset, limit int) []int64 {
+	if offset >= len(v) {
+		return nil
+	}
+	end := offset + limit
+	if end > len(v) {
+		end = len(v)
+	}
+	return v[offset:end]
+}
+
+func equalInt64s(a, b []int64) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for i := range a {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+	return true
+}

pkg/pgserver/connection.go

@@ -994,7 +994,13 @@ func (c *Connection) sendBackendKeyData(processID, secretKey int32) error {
 func (c *Connection) sendReadyForQuery() error {
 	mb := NewMessageBuilder()
 	mb.AppendByte(c.txStatus)
-	return c.writeMessage(MsgReadyForQuery, mb.Bytes())
+	if err := c.writeMessage(MsgReadyForQuery, mb.Bytes()); err != nil {
+		return err
+	}
+	// ReadyForQuery closes out a response cycle, so flush everything buffered
+	// so far. This is what makes simple-query results and Sync responses
+	// visible to the client.
+	return c.writer.Flush()
 }
 
 // sendEmptyQueryResponse sends empty query response
@@ -1062,18 +1068,25 @@ func (c *Connection) sendError(severity, code, message string) error {
 	mb.WriteString(message)
 	mb.AppendByte(0) // Terminator
 
-	return c.writeMessage(MsgErrorResponse, mb.Bytes())
+	if err := c.writeMessage(MsgErrorResponse, mb.Bytes()); err != nil {
+		return err
+	}
+	// Errors must be visible promptly, including FATAL startup failures that
+	// are not followed by a ReadyForQuery before the connection closes.
+	return c.writer.Flush()
 }
 
-// writeMessage writes a message to the connection
+// writeMessage buffers a message for the connection. It does not flush, so
+// callers that need to make a response visible to the client must flush at the
+// appropriate protocol boundary (ReadyForQuery, an explicit Flush message, or
+// an error response). Buffering amortizes the per-message syscalls that a
+// result set would otherwise incur; the underlying bufio.Writer bounds memory
+// use so large result sets cannot grow the buffer without limit.
 func (c *Connection) writeMessage(msgType byte, data []byte) error {
 	if !c.quiet {
 		log.Printf("Sending message type=%c length=%d", msgType, len(data)+4)
 	}
-	if err := WriteMessage(c.writer, msgType, data); err != nil {
-		return err
-	}
-	return c.writer.Flush()
+	return WriteMessage(c.writer, msgType, data)
 }
 
 // getOIDForType returns PostgreSQL OID for type

pkg/pgserver/connection_test.go

@@ -3,8 +3,11 @@ package pgserver
 import (
 	"bufio"
 	"bytes"
+	"io"
 	"net"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/danfragoso/pizzasql-next/pkg/executor"
 	"github.com/danfragoso/pizzasql-next/pkg/lexer"
@@ -230,3 +233,144 @@ func TestIsRollbackStatement(t *testing.T) {
 		}
 	}
 }
+
+// countingConn is an in-memory net.Conn that records how many times the
+// underlying stream is written. bufio.Writer flushes produce exactly one write
+// call each, so this counts flushes without blocking like net.Pipe does.
+type countingConn struct {
+	mu     sync.Mutex
+	buf    bytes.Buffer
+	writes int
+}
+
+func (c *countingConn) Read(p []byte) (int, error) { return 0, io.EOF }
+
+func (c *countingConn) Write(p []byte) (int, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.writes++
+	return c.buf.Write(p)
+}
+
+func (c *countingConn) Close() error                       { return nil }
+func (c *countingConn) LocalAddr() net.Addr                { return &net.TCPAddr{} }
+func (c *countingConn) RemoteAddr() net.Addr               { return &net.TCPAddr{} }
+func (c *countingConn) SetDeadline(t time.Time) error      { return nil }
+func (c *countingConn) SetReadDeadline(t time.Time) error  { return nil }
+func (c *countingConn) SetWriteDeadline(t time.Time) error { return nil }
+
+func (c *countingConn) bytes() []byte {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return append([]byte(nil), c.buf.Bytes()...)
+}
+
+func newCountingTestConnection(t *testing.T) (*Connection, *countingConn) {
+	t.Helper()
+	cc := &countingConn{}
+	c := &Connection{
+		conn:       cc,
+		reader:     bufio.NewReader(cc),
+		writer:     bufio.NewWriter(cc),
+		params:     map[string]string{"user": "tester"},
+		statements: make(map[string]*preparedStatement),
+		portals:    make(map[string]*portal),
+		txStatus:   TxStatusIdle,
+		quiet:      true,
+	}
+	return c, cc
+}
+
+func readAllMessages(t *testing.T, data []byte) []*Message {
+	t.Helper()
+	r := bytes.NewReader(data)
+	var msgs []*Message
+	for r.Len() > 0 {
+		m, err := ReadMessage(r)
+		if err != nil {
+			t.Fatalf("read message: %v", err)
+		}
+		msgs = append(msgs, m)
+	}
+	return msgs
+}
+
+func TestWriteMessageDoesNotFlush(t *testing.T) {
+	c, cc := newCountingTestConnection(t)
+	if err := c.writeMessage(MsgCommandComplete, []byte("SELECT 0")); err != nil {
+		t.Fatal(err)
+	}
+	if cc.writes != 0 {
+		t.Fatalf("writeMessage flushed %d times, want 0", cc.writes)
+	}
+}
+
+func TestReadyForQueryFlushesBufferedMessages(t *testing.T) {
+	c, cc := newCountingTestConnection(t)
+	if err := c.writeMessage(MsgCommandComplete, []byte("SELECT 1")); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.sendReadyForQuery(); err != nil {
+		t.Fatal(err)
+	}
+	if cc.writes == 0 {
+		t.Fatal("sendReadyForQuery did not flush buffered messages")
+	}
+	msgs := readAllMessages(t, cc.bytes())
+	if len(msgs) != 2 {
+		t.Fatalf("got %d messages, want 2", len(msgs))
+	}
+	if msgs[0].Type != MsgCommandComplete || msgs[1].Type != MsgReadyForQuery {
+		t.Fatalf("unexpected message sequence: %c, %c", msgs[0].Type, msgs[1].Type)
+	}
+}
+
+func TestErrorResponseFlushes(t *testing.T) {
+	c, cc := newCountingTestConnection(t)
+	if err := c.sendError("ERROR", ErrCodeSyntaxError, "boom"); err != nil {
+		t.Fatal(err)
+	}
+	if cc.writes == 0 {
+		t.Fatal("sendError did not flush")
+	}
+	msgs := readAllMessages(t, cc.bytes())
+	if len(msgs) != 1 || msgs[0].Type != MsgErrorResponse {
+		t.Fatalf("expected a single error response, got %d messages", len(msgs))
+	}
+}
+
+func TestFlushMessageFlushes(t *testing.T) {
+	c, cc := newCountingTestConnection(t)
+	if err := c.writeMessage(MsgCommandComplete, []byte("SELECT 1")); err != nil {
+		t.Fatal(err)
+	}
+	if err := c.handleMessage(&Message{Type: MsgFlush}); err != nil {
+		t.Fatal(err)
+	}
+	if cc.writes == 0 {
+		t.Fatal("Flush message did not flush buffered data")
+	}
+}
+
+func TestMultiRowResultBuffersRows(t *testing.T) {
+	c, cc := newCountingTestConnection(t)
+	result := executor.NewResult("SELECT")
+	result.AddColumnWithType("n", "INTEGER")
+	const rows = 500
+	for i := 0; i < rows; i++ {
+		result.AddRow(int64(i))
+	}
+	if err := c.sendResult(result, &parser.SelectStmt{}); err != nil {
+		t.Fatal(err)
+	}
+	// sendResult emits RowDescription + rows DataRows + CommandComplete. With
+	// per-message flushing that would be rows+2 underlying writes; buffered it
+	// is bounded by the bufio.Writer capacity (a few flushes at most).
+	numMessages := rows + 2
+	if cc.writes >= numMessages {
+		t.Fatalf("sendResult flushed %d times, want fewer than %d messages", cc.writes, numMessages)
+	}
+	if cc.writes >= rows {
+		t.Fatalf("sendResult flushed %d times for %d rows, expected buffering", cc.writes, rows)
+	}
+}

pkg/storage/count_cache_test.go

@@ -0,0 +1,432 @@
+package storage
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+)
+
+// TestCountFastLifecycle verifies the COUNT(*) metadata counter stays exact
+// across insert, update, delete, bulk insert, and truncate.
+func TestCountFastLifecycle(t *testing.T) {
+	kv := newTestKVServer(t)
+	defer kv.close()
+
+	pool := newTestKVPool(kv, 2, 5*time.Second)
+	defer pool.Close()
+
+	schemas := NewSchemaManager(pool, "testdb")
+	tables := NewTableManager(pool, schemas, "testdb")
+
+	err := schemas.CreateTable(&Schema{
+		Name: "users",
+		Columns: []Column{
+			{Name: "id", Type: "INTEGER", Nullable: false, PrimaryKey: true},
+			{Name: "name", Type: "TEXT", Nullable: true},
+		},
+	})
+	if err != nil {
+		t.Fatalf("create table: %v", err)
+	}
+
+	assertCount := func(want int) {
+		t.Helper()
+		got, err := tables.CountFast("users")
+		if err != nil {
+			t.Fatalf("CountFast: %v", err)
+		}
+		if got != want {
+			t.Fatalf("CountFast = %d, want %d", got, want)
+		}
+	}
+
+	assertCount(0)
+
+	for i := int64(1); i <= 3; i++ {
+		if err := tables.Insert("users", Row{"id": i, "name": fmt.Sprintf("u%d", i)}); err != nil {
+			t.Fatalf("insert %d: %v", i, err)
+		}
+	}
+	assertCount(3)
+
+	// UPDATE does not change the count.
+	n, err := tables.Update("users", Row{"name": "renamed"}, func(r Row) bool {
+		return fmt.Sprintf("%v", r["id"]) == "1"
+	})
+	if err != nil || n != 1 {
+		t.Fatalf("update: n=%d err=%v", n, err)
+	}
+	assertCount(3)
+
+	// DELETE decrements.
+	n, err = tables.Delete("users", func(r Row) bool {
+		return fmt.Sprintf("%v", r["id"]) == "2"
+	})
+	if err != nil || n != 1 {
+		t.Fatalf("delete: n=%d err=%v", n, err)
+	}
+	assertCount(2)
+
+	// Bulk insert adds len(rows).
+	n, err = tables.InsertBulk("users", []Row{
+		{"id": 4, "name": "u4"},
+		{"id": 5, "name": "u5"},
+		{"id": 6, "name": "u6"},
+	})
+	if err != nil || n != 3 {
+		t.Fatalf("bulk insert: n=%d err=%v", n, err)
+	}
+	assertCount(5)
+
+	// Truncate zeroes the count.
+	if _, err := tables.Truncate("users"); err != nil {
+		t.Fatalf("truncate: %v", err)
+	}
+	assertCount(0)
+}
+
+// TestCountFastDerivedAfterRestart verifies the counter is re-derived from
+// durable rows when a fresh TableManager (process restart) has no in-memory
+// count yet.
+func TestCountFastDerivedAfterRestart(t *testing.T) {
+	kv := newTestKVServer(t)
+	defer kv.close()
+
+	pool := newTestKVPool(kv, 2, 5*time.Second)
+	defer pool.Close()
+
+	schemas := NewSchemaManager(pool, "testdb")
+	tables := NewTableManager(pool, schemas, "testdb")
+
+	err := schemas.CreateTable(&Schema{
+		Name: "events",
+		Columns: []Column{
+			{Name: "id", Type: "INTEGER", Nullable: false, PrimaryKey: true},
+		},
+	})
+	if err != nil {
+		t.Fatalf("create table: %v", err)
+	}
+	for i := int64(1); i <= 4; i++ {
+		if err := tables.Insert("events", Row{"id": i}); err != nil {
+			t.Fatalf("insert %d: %v", i, err)
+		}
+	}
+
+	restartedSchemas := NewSchemaManager(pool, "testdb")
+	restartedTables := NewTableManager(pool, restartedSchemas, "testdb")
+
+	got, err := restartedTables.CountFast("events")
+	if err != nil {
+		t.Fatalf("CountFast after restart: %v", err)
+	}
+	if got != 4 {
+		t.Fatalf("CountFast after restart = %d, want 4", got)
+	}
+}
+
+// TestIncrementalCacheAndIndexMaintenance verifies that writes keep the row
+// cache and already-built in-memory indexes coherent without full-table cache
+// invalidation.
+func TestIncrementalCacheAndIndexMaintenance(t *testing.T) {
+	kv := newTestKVServer(t)
+	defer kv.close()
+
+	pool := newTestKVPool(kv, 2, 5*time.Second)
+	defer pool.Close()
+
+	schemas := NewSchemaManager(pool, "testdb")
+	tables := NewTableManager(pool, schemas, "testdb")
+
+	err := schemas.CreateTable(&Schema{
+		Name: "users",
+		Columns: []Column{
+			{Name: "id", Type: "INTEGER", Nullable: false, PrimaryKey: true},
+			{Name: "status", Type: "TEXT", Nullable: false},
+		},
+	})
+	if err != nil {
+		t.Fatalf("create table: %v", err)
+	}
+	err = schemas.CreateIndex(&Index{
+		Name:  "idx_users_status",
+		Table: "users",
+		Columns: []IndexColumn{
+			{Name: "status"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("create index: %v", err)
+	}
+
+	// Seed: 1 active, 1 inactive.
+	if err := tables.Insert("users", Row{"id": 1, "status": "active"}); err != nil {
+		t.Fatalf("insert 1: %v", err)
+	}
+	if err := tables.Insert("users", Row{"id": 2, "status": "inactive"}); err != nil {
+		t.Fatalf("insert 2: %v", err)
+	}
+
+	// Load the row cache and the index.
+	all, err := tables.Select("users", nil)
+	if err != nil || len(all) != 2 {
+		t.Fatalf("initial select: len=%d err=%v", len(all), err)
+	}
+	active, err := tables.SelectByIndex("users", "idx_users_status", "active")
+	if err != nil || len(active) != 1 {
+		t.Fatalf("initial active index: len=%d err=%v", len(active), err)
+	}
+
+	// Bulk insert two more active rows; the index must reflect them.
+	if n, err := tables.InsertBulk("users", []Row{
+		{"id": 3, "status": "active"},
+		{"id": 4, "status": "active"},
+	}); err != nil || n != 2 {
+		t.Fatalf("bulk insert: n=%d err=%v", n, err)
+	}
+	active, err = tables.SelectByIndex("users", "idx_users_status", "active")
+	if err != nil || len(active) != 3 {
+		t.Fatalf("active index after bulk: len=%d err=%v", len(active), err)
+	}
+	all, err = tables.Select("users", nil)
+	if err != nil || len(all) != 4 {
+		t.Fatalf("select after bulk: len=%d err=%v", len(all), err)
+	}
+
+	// Update one active -> inactive; both index buckets must stay coherent.
+	if n, err := tables.Update("users", Row{"status": "inactive"}, func(r Row) bool {
+		return fmt.Sprintf("%v", r["id"]) == "3"
+	}); err != nil || n != 1 {
+		t.Fatalf("update: n=%d err=%v", n, err)
+	}
+	active, _ = tables.SelectByIndex("users", "idx_users_status", "active")
+	inactive, _ := tables.SelectByIndex("users", "idx_users_status", "inactive")
+	if len(active) != 2 || len(inactive) != 2 {
+		t.Fatalf("indexes after update: active=%d inactive=%d", len(active), len(inactive))
+	}
+
+	// Delete one row; cache and index must shrink.
+	if n, err := tables.Delete("users", func(r Row) bool {
+		return fmt.Sprintf("%v", r["id"]) == "4"
+	}); err != nil || n != 1 {
+		t.Fatalf("delete: n=%d err=%v", n, err)
+	}
+	all, err = tables.Select("users", nil)
+	if err != nil || len(all) != 3 {
+		t.Fatalf("select after delete: len=%d err=%v", len(all), err)
+	}
+	active, _ = tables.SelectByIndex("users", "idx_users_status", "active")
+	if len(active) != 1 {
+		t.Fatalf("active index after delete: len=%d", len(active))
+	}
+}
+
+func TestClearIndexDropsInMemoryEntries(t *testing.T) {
+	kv := newTestKVServer(t)
+	defer kv.close()
+
+	pool := newTestKVPool(kv, 2, 5*time.Second)
+	defer pool.Close()
+
+	schemas := NewSchemaManager(pool, "testdb")
+	tables := NewTableManager(pool, schemas, "testdb")
+	if err := schemas.CreateTable(&Schema{
+		Name: "users",
+		Columns: []Column{
+			{Name: "id", Type: "INTEGER", Nullable: false, PrimaryKey: true},
+			{Name: "status", Type: "TEXT", Nullable: false},
+		},
+	}); err != nil {
+		t.Fatalf("create table: %v", err)
+	}
+	if err := schemas.CreateIndex(&Index{Name: "idx_status", Table: "users", Columns: []IndexColumn{{Name: "status"}}}); err != nil {
+		t.Fatalf("create index: %v", err)
+	}
+	if err := tables.Insert("users", Row{"id": 1, "status": "old"}); err != nil {
+		t.Fatalf("insert: %v", err)
+	}
+	if rows, err := tables.SelectByIndex("users", "idx_status", "old"); err != nil || len(rows) != 1 {
+		t.Fatalf("build index: len=%d err=%v", len(rows), err)
+	}
+	if err := tables.ClearIndex("idx_status", "users", []string{"status"}); err != nil {
+		t.Fatalf("clear index: %v", err)
+	}
+
+	tables.cacheMu.RLock()
+	_, cached := tables.indexCache["idx_status"]
+	_, mapped := tables.indexTable["idx_status"]
+	tables.cacheMu.RUnlock()
+	if cached || mapped {
+		t.Fatalf("cleared index remains cached: cache=%v table=%v", cached, mapped)
+	}
+	if rows, err := tables.SelectByIndex("users", "idx_status", "old"); err != nil || len(rows) != 0 {
+		t.Fatalf("cleared index rebuilt during drop: len=%d err=%v", len(rows), err)
+	}
+
+	if err := tables.BuildIndex("idx_status", "users", []string{"status"}); err != nil {
+		t.Fatalf("rebuild index: %v", err)
+	}
+	if rows, err := tables.SelectByIndex("users", "idx_status", "old"); err != nil || len(rows) != 1 {
+		t.Fatalf("rebuilt index: len=%d err=%v", len(rows), err)
+	}
+}
+
+func TestConcurrentDuplicateInsertKeepsCountExact(t *testing.T) {
+	kv := newTestKVServer(t)
+	defer kv.close()
+
+	pool := newTestKVPool(kv, 2, 5*time.Second)
+	defer pool.Close()
+
+	schemas := NewSchemaManager(pool, "testdb")
+	tables := NewTableManager(pool, schemas, "testdb")
+	if err := schemas.CreateTable(&Schema{
+		Name:    "users",
+		Columns: []Column{{Name: "id", Type: "INTEGER", Nullable: false, PrimaryKey: true}},
+	}); err != nil {
+		t.Fatalf("create table: %v", err)
+	}
+	if got, err := tables.CountFast("users"); err != nil || got != 0 {
+		t.Fatalf("initial count=%d err=%v", got, err)
+	}
+
+	start := make(chan struct{})
+	errs := make(chan error, 2)
+	for i := 0; i < 2; i++ {
+		go func() {
+			<-start
+			errs <- tables.Insert("users", Row{"id": 1})
+		}()
+	}
+	close(start)
+	successes := 0
+	for i := 0; i < 2; i++ {
+		if err := <-errs; err == nil {
+			successes++
+		}
+	}
+	if successes != 1 {
+		t.Fatalf("successful inserts=%d, want 1", successes)
+	}
+	if got, err := tables.CountFast("users"); err != nil || got != 1 {
+		t.Fatalf("final count=%d err=%v, want 1", got, err)
+	}
+}
+
+// TestConcurrentFirstLoadAndInsert runs the first cache load (Select on an
+// unloaded table) concurrently with inserts, then verifies the cache ends up
+// coherent with durable rows. Run under -race to detect map/slice data races.
+func TestConcurrentFirstLoadAndInsert(t *testing.T) {
+	kv := newTestKVServer(t)
+	defer kv.close()
+
+	pool := newTestKVPool(kv, 16, 10*time.Second)
+	defer pool.Close()
+
+	schemas := NewSchemaManager(pool, "testdb")
+	tables := NewTableManager(pool, schemas, "testdb")
+
+	err := schemas.CreateTable(&Schema{
+		Name: "users",
+		Columns: []Column{
+			{Name: "id", Type: "INTEGER", Nullable: false, PrimaryKey: true},
+			{Name: "name", Type: "TEXT", Nullable: true},
+		},
+	})
+	if err != nil {
+		t.Fatalf("create table: %v", err)
+	}
+
+	const n = 200
+	start := make(chan struct{})
+	var wg sync.WaitGroup
+	for i := 0; i < n; i++ {
+		i := i
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			if i%2 == 0 {
+				if err := tables.Insert("users", Row{"id": int64(i + 1), "name": fmt.Sprintf("u%d", i)}); err != nil {
+					t.Errorf("insert %d: %v", i, err)
+					return
+				}
+			} else {
+				if _, err := tables.Select("users", nil); err != nil {
+					t.Errorf("select: %v", err)
+				}
+			}
+		}()
+	}
+	close(start)
+	wg.Wait()
+
+	rows, err := tables.Select("users", nil)
+	if err != nil {
+		t.Fatalf("final select: %v", err)
+	}
+	if want := n / 2; len(rows) != want {
+		t.Fatalf("final select = %d rows, want %d", len(rows), want)
+	}
+	if got, _ := tables.CountFast("users"); got != n/2 {
+		t.Fatalf("final count = %d, want %d", got, n/2)
+	}
+}
+
+// TestConcurrentFirstCountFastAndInsert runs first-time count derivation
+// concurrently with inserts, then verifies the count ends up exact.
+func TestConcurrentFirstCountFastAndInsert(t *testing.T) {
+	kv := newTestKVServer(t)
+	defer kv.close()
+
+	pool := newTestKVPool(kv, 16, 10*time.Second)
+	defer pool.Close()
+
+	schemas := NewSchemaManager(pool, "testdb")
+	tables := NewTableManager(pool, schemas, "testdb")
+
+	err := schemas.CreateTable(&Schema{
+		Name: "items",
+		Columns: []Column{
+			{Name: "id", Type: "INTEGER", Nullable: false, PrimaryKey: true},
+		},
+	})
+	if err != nil {
+		t.Fatalf("create table: %v", err)
+	}
+
+	const n = 200
+	start := make(chan struct{})
+	var wg sync.WaitGroup
+	for i := 0; i < n; i++ {
+		i := i
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			if i%2 == 0 {
+				if err := tables.Insert("items", Row{"id": int64(i + 1)}); err != nil {
+					t.Errorf("insert %d: %v", i, err)
+				}
+			} else {
+				if _, err := tables.CountFast("items"); err != nil {
+					t.Errorf("count: %v", err)
+				}
+			}
+		}()
+	}
+	close(start)
+	wg.Wait()
+
+	if got, _ := tables.CountFast("items"); got != n/2 {
+		t.Fatalf("final count = %d, want %d", got, n/2)
+	}
+	rows, err := tables.Select("items", nil)
+	if err != nil {
+		t.Fatalf("final select: %v", err)
+	}
+	if len(rows) != n/2 {
+		t.Fatalf("final select = %d rows, want %d", len(rows), n/2)
+	}
+}

pkg/storage/schema.go

@@ -244,8 +244,12 @@ func (m *SchemaManager) DropTable(name string) error {
 func (m *SchemaManager) GetSchema(name string) (*Schema, error) {
 	m.mu.RLock()
 	if schema, ok := m.cache[strings.ToLower(name)]; ok {
+		// Clone while still holding the read lock: the cached schema's
+		// NextRowID field is mutated under the write lock, so cloning outside
+		// the lock races with that mutation.
+		cloned := cloneSchema(schema)
 		m.mu.RUnlock()
-		return cloneSchema(schema), nil
+		return cloned, nil
 	}
 	m.mu.RUnlock()
 

pkg/storage/table.go

@@ -23,21 +23,54 @@ type TableManager struct {
 
 	indexCache map[string]map[string][]int64 // index name → indexed value → rowids
 	indexTable map[string]string             // index name → table name
+	// disabledIndexes prevents a concurrent lookup from rebuilding an index
+	// after DROP has cleared it but before the schema entry is removed.
+	disabledIndexes map[string]bool
+
+	// counts holds exact per-table row counts for the COUNT(*) fast path.
+	// It is derived lazily from durable rows on first use and maintained
+	// incrementally by Insert/InsertBulk/Delete thereafter.
+	counts     map[string]int
+	countsInit map[string]bool
+
+	// locks is a map of per-table mutexes used to serialize cache/count/index
+	// loading (KV scan + install) against writes to the same table, so a scan
+	// cannot miss or double-count a concurrent write. Operations on different
+	// tables proceed concurrently. locksMu guards only the map itself and is
+	// never held across I/O or row operations.
+	locksMu sync.Mutex
+	locks   map[string]*sync.Mutex
 }
 
 // NewTableManager creates a new table manager.
 func NewTableManager(pool *KVPool, schema *SchemaManager, database string) *TableManager {
 	return &TableManager{
-		pool:       pool,
-		schema:     schema,
-		database:   database,
-		rowCache:   make(map[string][]Row),
-		rowIDMap:   make(map[string]map[int64]Row),
-		indexCache: make(map[string]map[string][]int64),
-		indexTable: make(map[string]string),
+		pool:            pool,
+		schema:          schema,
+		database:        database,
+		rowCache:        make(map[string][]Row),
+		rowIDMap:        make(map[string]map[int64]Row),
+		indexCache:      make(map[string]map[string][]int64),
+		indexTable:      make(map[string]string),
+		disabledIndexes: make(map[string]bool),
+		counts:          make(map[string]int),
+		countsInit:      make(map[string]bool),
+		locks:           make(map[string]*sync.Mutex),
 	}
 }
 
+// tableLock returns the per-table mutex keyed by lowercase table name.
+func (m *TableManager) tableLock(key string) *sync.Mutex {
+	m.locksMu.Lock()
+	l, ok := m.locks[key]
+	if !ok {
+		l = &sync.Mutex{}
+		m.locks[key] = l
+	}
+	m.locksMu.Unlock()
+	return l
+}
+
 // invalidateCache removes a table's rows from the in-memory cache.
 func (m *TableManager) invalidateCache(table string) {
 	m.cacheMu.Lock()
@@ -58,6 +91,186 @@ func (m *TableManager) InvalidateCache(table string) {
 	m.invalidateCache(table)
 }
 
+// loadTableLocked ensures the row cache for key is populated from durable rows.
+// The caller must hold the table's per-table lock so a concurrent write cannot
+// slip between the KV scan and the cache install.
+func (m *TableManager) loadTableLocked(key, table string) error {
+	m.cacheMu.RLock()
+	_, ok := m.rowCache[key]
+	m.cacheMu.RUnlock()
+	if ok {
+		return nil
+	}
+
+	prefix := m.dataPrefix(table)
+	var values []string
+	err := m.pool.WithClient(func(c *KVClient) error {
+		var err error
+		values, err = c.Reads(prefix)
+		return err
+	})
+	if err != nil {
+		return err
+	}
+
+	loaded := make([]Row, 0, len(values))
+	byRowID := make(map[int64]Row, len(values))
+	for _, data := range values {
+		var row Row
+		if err := json.Unmarshal([]byte(data), &row); err != nil {
+			continue
+		}
+		loaded = append(loaded, row)
+		if rowid, ok := valueAsInt64(row["_rowid_"]); ok {
+			byRowID[rowid] = row
+		}
+	}
+
+	m.cacheMu.Lock()
+	if _, ok := m.rowCache[key]; !ok {
+		m.rowCache[key] = loaded
+		m.rowIDMap[key] = byRowID
+	}
+	m.cacheMu.Unlock()
+
+	return nil
+}
+
+// loadTable populates the row cache for table, acquiring the per-table lock.
+func (m *TableManager) loadTable(key, table string) error {
+	tl := m.tableLock(key)
+	tl.Lock()
+	defer tl.Unlock()
+	return m.loadTableLocked(key, table)
+}
+
+// CountFast returns the exact number of rows in a table. The count is derived
+// from durable rows on first use (recovering across restarts) and then
+// maintained incrementally by the write paths, so repeated COUNT(*) queries
+// avoid a full table scan. It intentionally does not persist a counter to KV:
+// the KV layer has no atomic increment primitive, and a durable counter that
+// could diverge from the rows on crash would be worse than a lazily-derived,
+// always-exact value. The cost is one table scan the first time COUNT(*) is
+// issued after startup.
+func (m *TableManager) CountFast(table string) (int, error) {
+	key := strings.ToLower(table)
+
+	m.cacheMu.RLock()
+	init := m.countsInit[key]
+	n := m.counts[key]
+	m.cacheMu.RUnlock()
+	if init {
+		return n, nil
+	}
+
+	// Serialize first-time derivation against writes to this table so a
+	// concurrent insert/delete cannot be missed or double-counted.
+	tl := m.tableLock(key)
+	tl.Lock()
+	defer tl.Unlock()
+
+	m.cacheMu.RLock()
+	init = m.countsInit[key]
+	n = m.counts[key]
+	m.cacheMu.RUnlock()
+	if init {
+		return n, nil
+	}
+
+	prefix := m.dataPrefix(table)
+	var values []string
+	err := m.pool.WithClient(func(c *KVClient) error {
+		var err error
+		values, err = c.Reads(prefix)
+		return err
+	})
+	if err != nil {
+		return 0, err
+	}
+
+	m.cacheMu.Lock()
+	m.counts[key] = len(values)
+	m.countsInit[key] = true
+	m.cacheMu.Unlock()
+
+	return len(values), nil
+}
+
+// incrCount adjusts the derived per-table row count. It is a no-op until the
+// count has been initialized, since an uninitialized count is re-derived from
+// durable rows (which already reflect the write) on next use.
+func (m *TableManager) incrCount(table string, delta int) {
+	key := strings.ToLower(table)
+	m.cacheMu.Lock()
+	if m.countsInit[key] {
+		m.counts[key] += delta
+	}
+	m.cacheMu.Unlock()
+}
+
+// cacheInsert adds a row to the in-memory row cache if it is already loaded.
+// It is idempotent: a rowid already present is not appended twice, so a
+// partially-observed bulk insert cannot duplicate cache entries.
+func (m *TableManager) cacheInsert(table string, row Row) {
+	key := strings.ToLower(table)
+	rowid, ok := rowIDFromRow(row)
+	m.cacheMu.Lock()
+	defer m.cacheMu.Unlock()
+	byRowID, loaded := m.rowIDMap[key]
+	if !loaded {
+		return
+	}
+	if ok {
+		if _, exists := byRowID[rowid]; exists {
+			return
+		}
+		byRowID[rowid] = row
+	}
+	m.rowCache[key] = append(m.rowCache[key], row)
+}
+
+// cacheDelete removes a row from the in-memory row cache if it is already loaded.
+func (m *TableManager) cacheDelete(table string, row Row) {
+	key := strings.ToLower(table)
+	rowid, ok := rowIDFromRow(row)
+	m.cacheMu.Lock()
+	defer m.cacheMu.Unlock()
+	if ok {
+		if byRowID, exists := m.rowIDMap[key]; exists {
+			delete(byRowID, rowid)
+		}
+	}
+	if cached, exists := m.rowCache[key]; exists && ok {
+		for i, r := range cached {
+			if rid, rok := rowIDFromRow(r); rok && rid == rowid {
+				m.rowCache[key] = append(cached[:i], cached[i+1:]...)
+				break
+			}
+		}
+	}
+}
+
+// cacheUpdate replaces a row in the in-memory row cache if it is already loaded.
+func (m *TableManager) cacheUpdate(table string, row Row) {
+	key := strings.ToLower(table)
+	rowid, ok := rowIDFromRow(row)
+	m.cacheMu.Lock()
+	defer m.cacheMu.Unlock()
+	if ok {
+		if byRowID, exists := m.rowIDMap[key]; exists {
+			byRowID[rowid] = row
+		}
+	}
+	if cached, exists := m.rowCache[key]; exists && ok {
+		for i, r := range cached {
+			if rid, rok := rowIDFromRow(r); rok && rid == rowid {
+				m.rowCache[key][i] = row
+				break
+			}
+		}
+	}
+}
+
 // dataKey returns the key for a row.
 func (m *TableManager) dataKey(table, pk string) string {
 	return fmt.Sprintf("%s:_data:%s:%s", m.database, strings.ToLower(table), pk)
@@ -125,8 +338,13 @@ func (m *TableManager) Insert(table string, row Row) error {
 	}
 
 	pk := fmt.Sprintf("%v", pkValue)
+	tl := m.tableLock(strings.ToLower(table))
+	tl.Lock()
+	defer tl.Unlock()
 
-	// Check for duplicate
+	// Keep the duplicate check and write in one per-table critical section so
+	// concurrent inserts of the same primary key cannot both update the cache
+	// and row count for a single durable row.
 	key := m.dataKey(table, pk)
 	err = m.pool.WithClient(func(c *KVClient) error {
 		_, err := c.Read(key)
@@ -189,7 +407,6 @@ func (m *TableManager) Insert(table string, row Row) error {
 		return fmt.Errorf("failed to serialize row: %w", err)
 	}
 
-	// Write row
 	err = m.pool.WithClient(func(c *KVClient) error {
 		return c.Write(key, string(data))
 	})
@@ -200,7 +417,8 @@ func (m *TableManager) Insert(table string, row Row) error {
 	// Update in-memory indexes only. Durable index entries are derived from rows.
 	m.updateIndexesForRow(table, normalizedRow, true)
 
-	m.invalidateCache(table)
+	m.cacheInsert(table, normalizedRow)
+	m.incrCount(table, 1)
 	return nil
 }
 
@@ -283,6 +501,12 @@ func (m *TableManager) InsertBulk(table string, rows []Row) (int, error) {
 		rowKVs = append(rowKVs, kv{m.dataKey(table, pk), string(data)})
 	}
 
+	// Hold the per-table lock for the whole write+maintain phase so a
+	// concurrent cache/count load cannot scan a partially-written table.
+	tl := m.tableLock(strings.ToLower(table))
+	tl.Lock()
+	defer tl.Unlock()
+
 	// Write rows concurrently.
 	errs := make([]error, len(rowKVs))
 	var wg sync.WaitGroup
@@ -297,14 +521,25 @@ func (m *TableManager) InsertBulk(table string, rows []Row) (int, error) {
 		}()
 	}
 	wg.Wait()
-	for _, e := range errs {
+
+	// Maintain in-memory caches only for rows that actually persisted, so a
+	// partial failure cannot leave an already-loaded cache/count stale.
+	var firstErr error
+	numOK := 0
+	for i, e := range errs {
 		if e != nil {
-			return 0, e
+			if firstErr == nil {
+				firstErr = e
+			}
+			continue
 		}
+		m.updateIndexesForRow(table, normalized[i], true)
+		m.cacheInsert(table, normalized[i])
+		numOK++
 	}
+	m.incrCount(table, numOK)
 
-	m.invalidateCache(table)
-	return len(normalized), nil
+	return numOK, firstErr
 }
 
 // updateIndexesForRow adds or removes entries from already-built in-memory
@@ -351,54 +586,28 @@ func (m *TableManager) Select(table string, filter func(Row) bool) ([]Row, error
 	}
 
 	key := strings.ToLower(table)
+	if err := m.loadTable(key, table); err != nil {
+		return nil, err
+	}
 
+	// Snapshot row references under the read lock, then filter and clone only
+	// matching rows without holding a lock. Published cached rows are immutable:
+	// writers replace row references rather than mutating their maps in place.
+	// This keeps selective scans from allocating a map for every examined row.
 	m.cacheMu.RLock()
-	cached, ok := m.rowCache[key]
+	cached := m.rowCache[key]
+	snapshot := append([]Row(nil), cached...)
 	m.cacheMu.RUnlock()
 
-	if !ok {
-		prefix := m.dataPrefix(table)
-		var values []string
-		err := m.pool.WithClient(func(c *KVClient) error {
-			var err error
-			values, err = c.Reads(prefix)
-			return err
-		})
-		if err != nil {
-			return nil, err
-		}
-
-		loaded := make([]Row, 0, len(values))
-		byRowID := make(map[int64]Row, len(values))
-		for _, data := range values {
-			var row Row
-			if err := json.Unmarshal([]byte(data), &row); err != nil {
-				continue
-			}
-			loaded = append(loaded, row)
-			if rowid, ok := valueAsInt64(row["_rowid_"]); ok {
-				byRowID[rowid] = row
-			}
-		}
-
-		m.cacheMu.Lock()
-		m.rowCache[key] = loaded
-		m.rowIDMap[key] = byRowID
-		m.cacheMu.Unlock()
-
-		cached = loaded
-	}
-
 	if filter == nil {
-		result := make([]Row, len(cached))
-		for i, row := range cached {
-			result[i] = cloneRow(row)
+		rows := make([]Row, len(snapshot))
+		for i, row := range snapshot {
+			rows[i] = cloneRow(row)
 		}
-		return result, nil
+		return rows, nil
 	}
-
-	rows := make([]Row, 0, len(cached))
-	for _, row := range cached {
+	rows := make([]Row, 0, len(snapshot))
+	for _, row := range snapshot {
 		if filter(row) {
 			rows = append(rows, cloneRow(row))
 		}
@@ -453,9 +662,15 @@ func (m *TableManager) Update(table string, updates Row, filter func(Row) bool)
 		return 0, err
 	}
 
+	tl := m.tableLock(strings.ToLower(table))
+	tl.Lock()
+	defer tl.Unlock()
+
 	count := 0
 	for _, row := range rows {
-		// Remove old index entries before update
+		// Snapshot the pre-update row so removed index entries can be restored
+		// if persistence fails.
+		oldRow := cloneRow(row)
 		m.updateIndexesForRow(table, row, false)
 
 		// Apply updates
@@ -476,6 +691,7 @@ func (m *TableManager) Update(table string, updates Row, filter func(Row) bool)
 		// Serialize row
 		data, err := json.Marshal(row)
 		if err != nil {
+			m.updateIndexesForRow(table, oldRow, true)
 			continue
 		}
 
@@ -487,11 +703,13 @@ func (m *TableManager) Update(table string, updates Row, filter func(Row) bool)
 		if err == nil {
 			// Add new index entries after update
 			m.updateIndexesForRow(table, row, true)
+			m.cacheUpdate(table, row)
 			count++
+		} else {
+			m.updateIndexesForRow(table, oldRow, true)
 		}
 	}
 
-	m.invalidateCache(table)
 	return count, nil
 }
 
@@ -509,14 +727,19 @@ func (m *TableManager) UpdateFunc(table string, updateFn func(Row) (Row, error),
 		return 0, err
 	}
 
+	tl := m.tableLock(strings.ToLower(table))
+	tl.Lock()
+	defer tl.Unlock()
+
 	count := 0
 	for _, row := range rows {
-		// Remove old index entries before update
+		oldRow := cloneRow(row)
 		m.updateIndexesForRow(table, row, false)
 
 		// Compute updates using the provided function
 		updates, err := updateFn(row)
 		if err != nil {
+			m.updateIndexesForRow(table, oldRow, true)
 			return count, err
 		}
 
@@ -538,6 +761,7 @@ func (m *TableManager) UpdateFunc(table string, updateFn func(Row) (Row, error),
 		// Serialize row
 		data, err := json.Marshal(row)
 		if err != nil {
+			m.updateIndexesForRow(table, oldRow, true)
 			continue
 		}
 
@@ -549,11 +773,13 @@ func (m *TableManager) UpdateFunc(table string, updateFn func(Row) (Row, error),
 		if err == nil {
 			// Add new index entries after update
 			m.updateIndexesForRow(table, row, true)
+			m.cacheUpdate(table, row)
 			count++
+		} else {
+			m.updateIndexesForRow(table, oldRow, true)
 		}
 	}
 
-	m.invalidateCache(table)
 	return count, nil
 }
 
@@ -570,6 +796,10 @@ func (m *TableManager) Delete(table string, filter func(Row) bool) (int, error)
 		return 0, err
 	}
 
+	tl := m.tableLock(strings.ToLower(table))
+	tl.Lock()
+	defer tl.Unlock()
+
 	count := 0
 	for _, row := range rows {
 		// Remove index entries before deleting row
@@ -583,11 +813,15 @@ func (m *TableManager) Delete(table string, filter func(Row) bool) (int, error)
 			return c.Delete(key)
 		})
 		if err == nil {
+			m.cacheDelete(table, row)
 			count++
+		} else {
+			// Restore the index entries removed above.
+			m.updateIndexesForRow(table, row, true)
 		}
 	}
 
-	m.invalidateCache(table)
+	m.incrCount(table, -count)
 	return count, nil
 }
 
@@ -695,22 +929,46 @@ func rowIDFromRow(row Row) (int64, bool) {
 func (m *TableManager) ensureIndex(index *Index) error {
 	indexKey := strings.ToLower(index.Name)
 	m.cacheMu.RLock()
+	disabled := m.disabledIndexes[indexKey]
 	_, initialized := m.indexCache[indexKey]
 	m.cacheMu.RUnlock()
+	if disabled {
+		return nil
+	}
 	if initialized {
 		return nil
 	}
 
-	columns := make([]string, len(index.Columns))
-	for i, col := range index.Columns {
-		columns[i] = col.Name
+	// Serialize index build against writes to the same table so the derived
+	// entries cannot miss a concurrently-inserted row.
+	table := index.Table
+	key := strings.ToLower(table)
+	tl := m.tableLock(key)
+	tl.Lock()
+	defer tl.Unlock()
+
+	m.cacheMu.RLock()
+	disabled = m.disabledIndexes[indexKey]
+	_, initialized = m.indexCache[indexKey]
+	m.cacheMu.RUnlock()
+	if disabled {
+		return nil
+	}
+	if initialized {
+		return nil
 	}
 
-	rows, err := m.Select(index.Table, nil)
-	if err != nil {
+	if err := m.loadTableLocked(key, table); err != nil {
 		return err
 	}
 
+	columns := make([]string, len(index.Columns))
+	for i, col := range index.Columns {
+		columns[i] = col.Name
+	}
+
+	m.cacheMu.RLock()
+	rows := m.rowCache[key]
 	values := make(map[string][]int64)
 	for _, row := range rows {
 		rowid, ok := rowIDFromRow(row)
@@ -721,11 +979,12 @@ func (m *TableManager) ensureIndex(index *Index) error {
 		valueKey := formatIndexValue(colValue)
 		values[valueKey] = append(values[valueKey], rowid)
 	}
+	m.cacheMu.RUnlock()
 
 	m.cacheMu.Lock()
 	if _, initialized := m.indexCache[indexKey]; !initialized {
 		m.indexCache[indexKey] = values
-		m.indexTable[indexKey] = strings.ToLower(index.Table)
+		m.indexTable[indexKey] = key
 	}
 	m.cacheMu.Unlock()
 
@@ -805,6 +1064,13 @@ func (m *TableManager) LookupIndex(indexName string, colValue interface{}) ([]in
 
 // ClearIndex removes all entries for an index by scanning table and removing entries.
 func (m *TableManager) ClearIndex(indexName, tableName string, columns []string) error {
+	indexKey := strings.ToLower(indexName)
+	m.cacheMu.Lock()
+	delete(m.indexCache, indexKey)
+	delete(m.indexTable, indexKey)
+	m.disabledIndexes[indexKey] = true
+	m.cacheMu.Unlock()
+
 	rows, err := m.Select(tableName, nil)
 	if err != nil {
 		return err
@@ -823,6 +1089,13 @@ func (m *TableManager) ClearIndex(indexName, tableName string, columns []string)
 
 // BuildIndex builds index entries for all existing rows in a table.
 func (m *TableManager) BuildIndex(indexName, tableName string, columns []string) error {
+	indexKey := strings.ToLower(indexName)
+	m.cacheMu.Lock()
+	delete(m.disabledIndexes, indexKey)
+	delete(m.indexCache, indexKey)
+	delete(m.indexTable, indexKey)
+	m.cacheMu.Unlock()
+
 	index, err := m.schema.GetIndex(indexName)
 	if err == nil {
 		return m.ensureIndex(index)
@@ -844,7 +1117,6 @@ func (m *TableManager) BuildIndex(indexName, tableName string, columns []string)
 		values[formatIndexValue(colValue)] = append(values[formatIndexValue(colValue)], rowid)
 	}
 
-	indexKey := strings.ToLower(indexName)
 	m.cacheMu.Lock()
 	m.indexCache[indexKey] = values
 	m.indexTable[indexKey] = strings.ToLower(tableName)
@@ -896,20 +1168,15 @@ func (m *TableManager) SelectByIndex(table, indexName string, colValue interface
 		return []Row{}, nil
 	}
 
-	// Build a set of target rowids for O(1) lookup.
+	// Ensure the rowID map is loaded, then look up and clone rows under the
+	// read lock so writers cannot mutate the map concurrently.
 	key := strings.ToLower(table)
-	m.cacheMu.RLock()
-	byRowID, ok := m.rowIDMap[key]
-	m.cacheMu.RUnlock()
-	if !ok {
-		if _, err := m.Select(table, nil); err != nil {
-			return nil, err
-		}
-		m.cacheMu.RLock()
-		byRowID = m.rowIDMap[key]
-		m.cacheMu.RUnlock()
+	if err := m.loadTable(key, table); err != nil {
+		return nil, err
 	}
 
+	m.cacheMu.RLock()
+	byRowID := m.rowIDMap[key]
 	rows := make([]Row, 0, len(rowids))
 	seen := make(map[int64]struct{}, len(rowids))
 	for _, rid := range rowids {
@@ -918,9 +1185,10 @@ func (m *TableManager) SelectByIndex(table, indexName string, colValue interface
 		}
 		seen[rid] = struct{}{}
 		if row, ok := byRowID[rid]; ok {
-			rows = append(rows, row)
+			rows = append(rows, cloneRow(row))
 		}
 	}
+	m.cacheMu.RUnlock()
 
 	return rows, nil
 }