Forráskód Böngészése

fix: preserve predicates and table binding in equality lookups (#15)

Danilo Fragoso 3 napja
szülő
commit
0b6772069a

+ 55 - 7
pkg/executor/executor.go

@@ -514,13 +514,15 @@ func (e *Executor) executeSelect(stmt *parser.SelectStmt) (*Result, error) {
 			break
 		}
 	}
-	appliedInScan := effectiveWhere != nil && !whereHasSubquery && stmt.From[0].Alias == "" && !isMultiTable && stmt.From[0].Join == nil
+	canFilterInScan := effectiveWhere != nil && !whereHasSubquery && stmt.From[0].Alias == "" && !isMultiTable && stmt.From[0].Join == nil
+	appliedInScan := false
 
+	// A lookup may only restrict a correctly bound, preserved table. The
+	// original predicate is still evaluated after the lookup/join.
 	if effectiveWhere != nil && !isMultiTable {
-		// Check if we can use an index
-		colName, colValue, isEquality := e.extractIndexableCondition(stmt.Where)
+		colName, colValue, isEquality := e.extractSelectLookupCondition(stmt.Where, stmt.From[0])
 		if isEquality {
-			if strings.EqualFold(schema.PrimaryKey, colName) {
+			if strings.EqualFold(schema.PrimaryKey, colName) && safePrimaryKeyLiteral(colValue) {
 				row, getErr := e.session.GetByPK(tableName, fmt.Sprintf("%v", colValue))
 				if getErr == nil {
 					normalizeRowBySchema(row, schema)
@@ -555,7 +557,8 @@ func (e *Executor) executeSelect(stmt *parser.SelectStmt) (*Result, error) {
 	if !usedIndex {
 		var filterErr error
 		var filter func(storage.Row) bool
-		if appliedInScan {
+		if canFilterInScan {
+			appliedInScan = true
 			filter = func(row storage.Row) bool {
 				val, ferr := e.evalExpr(effectiveWhere, row)
 				if ferr != nil {
@@ -4333,7 +4336,7 @@ func (e *Executor) extractIndexableCondition(expr parser.Expr) (string, interfac
 	colRef, leftIsCol := binExpr.Left.(*parser.ColumnRef)
 	litExpr, rightIsLit := binExpr.Right.(*parser.LiteralExpr)
 
-	if leftIsCol && rightIsLit {
+	if leftIsCol && rightIsLit && colRef.Table == "" {
 		val, _ := e.evalLiteral(litExpr)
 		return colRef.Column, val, true
 	}
@@ -4342,7 +4345,7 @@ func (e *Executor) extractIndexableCondition(expr parser.Expr) (string, interfac
 	litExpr, leftIsLit := binExpr.Left.(*parser.LiteralExpr)
 	colRef, rightIsCol := binExpr.Right.(*parser.ColumnRef)
 
-	if leftIsLit && rightIsCol {
+	if leftIsLit && rightIsCol && colRef.Table == "" {
 		val, _ := e.evalLiteral(litExpr)
 		return colRef.Column, val, true
 	}
@@ -4350,6 +4353,51 @@ func (e *Executor) extractIndexableCondition(expr parser.Expr) (string, interfac
 	return "", nil, false
 }
 
+func (e *Executor) extractSelectLookupCondition(expr parser.Expr, table parser.TableRef) (string, interface{}, bool) {
+	for join := table.Join; join != nil; join = join.Table.Join {
+		if join.Type == parser.JoinRight || join.Type == parser.JoinFull {
+			return "", nil, false
+		}
+	}
+	bin, ok := expr.(*parser.BinaryExpr)
+	if !ok {
+		return "", nil, false
+	}
+	copyExpr := *bin
+	for _, side := range []*parser.Expr{&copyExpr.Left, &copyExpr.Right} {
+		if column, ok := (*side).(*parser.ColumnRef); ok {
+			if column.Table == "" && table.Join != nil {
+				return "", nil, false
+			}
+			qualifier := table.Alias
+			if qualifier == "" {
+				qualifier = table.Name
+			}
+			if column.Table != "" && !strings.EqualFold(column.Table, qualifier) {
+				return "", nil, false
+			}
+			copyColumn := *column
+			copyColumn.Table = ""
+			*side = &copyColumn
+		}
+	}
+	return e.extractIndexableCondition(&copyExpr)
+}
+
+// Numeric strings can compare equal to differently formatted primary keys;
+// a direct durable-key lookup would miss them. Keep those comparisons on the
+// scan path until primary-key affinity has a canonical storage representation.
+func safePrimaryKeyLiteral(value interface{}) bool {
+	if value == nil {
+		return false
+	}
+	if text, ok := value.(string); ok {
+		_, numeric := toNumeric(text)
+		return !numeric
+	}
+	return true
+}
+
 // executePragma executes a PRAGMA statement.
 func (e *Executor) executePragma(stmt *parser.PragmaStmt) (*Result, error) {
 	switch stmt.Name {

+ 57 - 0
pkg/executor/index_predicate_test.go

@@ -0,0 +1,57 @@
+package executor
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestIndexPredicateMatchesScan(t *testing.T) {
+	_, schema, table := newTestDB(t)
+	e := newExec(schema, table)
+	execMust(t, e, "CREATE TABLE t (id INTEGER PRIMARY KEY, v REAL, txt TEXT)")
+	execMust(t, e, "INSERT INTO t VALUES (1,1.0000001,'01'),(2,1.0000002,'1'),(3,NULL,'other')")
+	queries := []string{
+		"SELECT id FROM t WHERE v=1.0000001 ORDER BY id",
+		"SELECT id FROM t WHERE v=NULL ORDER BY id",
+		"SELECT id FROM t WHERE 1.0000002=v ORDER BY id",
+		"SELECT id FROM t WHERE txt='1' ORDER BY id",
+		"SELECT id FROM t WHERE id=NULL ORDER BY id",
+		"SELECT id FROM t AS x WHERE x.v=1.0000001 ORDER BY id",
+	}
+	want := make([][][]interface{}, len(queries))
+	for i, q := range queries {
+		want[i] = execMust(t, e, q).Rows
+	}
+	execMust(t, e, "CREATE INDEX v_idx ON t(v)")
+	execMust(t, e, "CREATE INDEX txt_idx ON t(txt)")
+	for _, tx := range []bool{false, true} {
+		if tx {
+			execMust(t, e, "BEGIN")
+		}
+		for i, q := range queries {
+			if got := execMust(t, e, q).Rows; !reflect.DeepEqual(got, want[i]) {
+				t.Errorf("tx=%v %s: got %v want %v", tx, q, got, want[i])
+			}
+		}
+		if tx {
+			execMust(t, e, "ROLLBACK")
+		}
+	}
+}
+
+func TestQualifiedJoinPredicateDoesNotRestrictLeftTable(t *testing.T) {
+	_, schema, table := newTestDB(t)
+	e := newExec(schema, table)
+	execMust(t, e, "CREATE TABLE lhs (id INTEGER PRIMARY KEY, x INTEGER)")
+	execMust(t, e, "CREATE TABLE rhs (id INTEGER PRIMARY KEY, x INTEGER)")
+	execMust(t, e, "INSERT INTO lhs VALUES (1,7),(2,7)")
+	execMust(t, e, "INSERT INTO rhs VALUES (1,7)")
+	for _, join := range []string{"JOIN", "LEFT JOIN"} {
+		q := "SELECT lhs.id,rhs.id FROM lhs " + join + " rhs ON lhs.x=rhs.x WHERE rhs.id=1 ORDER BY lhs.id"
+		got := execMust(t, e, q).Rows
+		want := [][]interface{}{{int64(1), int64(1)}, {int64(2), int64(1)}}
+		if !reflect.DeepEqual(got, want) {
+			t.Fatalf("%s: got %v want %v", q, got, want)
+		}
+	}
+}

+ 35 - 0
pkg/pgserver/index_predicate_test.go

@@ -0,0 +1,35 @@
+package pgserver
+
+import "testing"
+
+func TestIndexedEqualityPreservesSQLNullAndFloatSemantics(t *testing.T) {
+	c, client := newDataConnection(t)
+	for _, sql := range []string{
+		"CREATE TABLE t (id INTEGER PRIMARY KEY, v REAL)",
+		"INSERT INTO t VALUES (1,1.0000001),(2,1.0000002),(3,NULL)",
+		"CREATE INDEX t_v ON t(v)",
+	} {
+		runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte(sql), 0)})
+	}
+	for _, tc := range []struct {
+		sql  string
+		rows int
+	}{
+		{"SELECT id FROM t WHERE v=1.0000001", 1},
+		{"SELECT id FROM t WHERE v=NULL", 0},
+	} {
+		msgs := runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte(tc.sql), 0)})
+		count := 0
+		for _, msg := range msgs {
+			if msg.Type == MsgErrorResponse {
+				t.Fatal(errorMessage(msg))
+			}
+			if msg.Type == MsgDataRow {
+				count++
+			}
+		}
+		if count != tc.rows {
+			t.Fatalf("%s: got %d rows, want %d", tc.sql, count, tc.rows)
+		}
+	}
+}

+ 24 - 6
pkg/storage/table.go

@@ -1368,19 +1368,37 @@ func (m *TableManager) indexPrefix(indexName string) string {
 }
 
 func formatIndexValue(value interface{}) string {
+	// Equality predicates compare numeric values (including numeric text) as
+	// float64. Use the same equivalence classes without decimal rounding.
+	// Index hits are candidates: the executor still evaluates the predicate.
+	var number float64
 	switch v := value.(type) {
 	case float64:
-		if v == float64(int64(v)) {
-			return fmt.Sprintf("%d", int64(v))
-		}
-		return fmt.Sprintf("%f", v)
+		number = v
 	case int64:
-		return fmt.Sprintf("%d", v)
+		number = float64(v)
 	case int:
-		return fmt.Sprintf("%d", v)
+		number = float64(v)
+	case bool:
+		if v {
+			return "1"
+		}
+		return "0"
+	case []byte:
+		return formatIndexValue(string(v))
+	case string:
+		parsed, err := strconv.ParseFloat(v, 64)
+		if err != nil {
+			return v
+		}
+		number = parsed
 	default:
 		return fmt.Sprintf("%v", v)
 	}
+	if number == 0 {
+		return "0"
+	}
+	return strconv.FormatFloat(number, 'g', -1, 64)
 }
 
 func rowIDFromRow(row Row) (int64, bool) {