fix(disttae): validate partitioned ODKU against the resolved locked key

pkg/partitionprune/filter.go

@@ -148,7 +148,7 @@ func hashFilterExpr(
 			if !ok {
 				return nil, false, nil
 			}
-			if !matchesPartitionColumn(left.Col, colPosition, metadata.Partitions[0].Expr) {
+			if !matchesPartitionColumns(left.Col, colPosition, metadata) {
 				return nil, false, nil
 			}
 			value, canPrune := normalizePartitionValue(exprImpl.F.Args[1])
@@ -162,7 +162,11 @@ func hashFilterExpr(
 				exprs[i] = p.DeepCopyExpr(pt.Expr)
 			}
 			for i := range exprs {
-				mustReplaceCol(exprs[i], value)
+				var err error
+				exprs[i], err = mustReplaceCol(proc.Ctx, exprs[i], value)
+				if err != nil {
+					return nil, false, err
+				}
 			}
 			targets, err := filterResult(proc, exprs, metadata)
 			if err != nil {
@@ -263,7 +267,7 @@ func rangeFilterExpr(
 			if !ok {
 				return nil, false, nil
 			}
-			if !matchesPartitionColumn(left.Col, colPosition, metadata.Partitions[0].Expr) {
+			if !matchesPartitionColumns(left.Col, colPosition, metadata) {
 				return nil, false, nil
 			}
 			value, canPrune := normalizePartitionValue(exprImpl.F.Args[1])
@@ -279,7 +283,11 @@ func rangeFilterExpr(
 			for i := range exprs {
 				// a = 1 =>
 				// p1 <= 1 < p2
-				mustReplaceCol(exprs[i], value)
+				var err error
+				exprs[i], err = mustReplaceCol(proc.Ctx, exprs[i], value)
+				if err != nil {
+					return nil, false, err
+				}
 			}
 			targets, err := filterResult(proc, exprs, metadata)
 			if err != nil {
@@ -406,42 +414,44 @@ func normalizePartitionValue(value *plan.Expr) (*plan.Expr, bool) {
 	return convertFoldExprToNormal(p.DeepCopyExpr(value))
 }
 
-// mustReplaceCol replaces column references in an expression with a given value expression.
-// It recursively traverses the expression tree to find and replace column references.
-func mustReplaceCol(expr, value *plan.Expr) {
+// mustReplaceCol substitutes every reference after callers establish that all
+// partition expressions depend on the matched filter column. Cast each value
+// to the original reference type and propagate construction failures.
+func mustReplaceCol(ctx context.Context, expr, value *plan.Expr) (*plan.Expr, error) {
 	switch e := expr.Expr.(type) {
+	case *plan.Expr_Col:
+		return appendCastBeforeExpr(ctx, value, expr.Typ)
 	case *plan.Expr_F:
-		for i := range e.F.Args {
-			switch e.F.Args[i].Expr.(type) {
-			case *plan.Expr_Col:
-				e.F.Args[i], _ = appendCastBeforeExpr(context.Background(), value, e.F.Args[i].Typ)
-				return
-			case *plan.Expr_F:
-				mustReplaceCol(e.F.Args[i], value)
+		for i, arg := range e.F.Args {
+			replacement, err := mustReplaceCol(ctx, arg, value)
+			if err != nil {
+				return nil, err
 			}
+			e.F.Args[i] = replacement
 		}
 	}
+	return expr, nil
 }
 
-// mustReplaceColPos recursively traverses the expression tree to find and replace
-// the column position with the specified position. This is used in list partitioning
-// to reset column positions when evaluating expressions against constructed batches.
+// mustReplaceColPos remaps every occurrence of the first source column. Other
+// columns keep their schema positions. RowID offsets are already represented
+// by the metadata expression selected by the caller.
 func mustReplaceColPos(expr *plan.Expr, pos int32) {
-	switch e := expr.Expr.(type) {
-	case *plan.Expr_F:
-		for i := range e.F.Args {
-			switch col := e.F.Args[i].Expr.(type) {
-			case *plan.Expr_Col:
-				if pos == -1 {
-					pos = col.Col.ColPos + 1
-				}
-				col.Col.ColPos = pos
-				return
-			case *plan.Expr_F:
-				mustReplaceColPos(e.F.Args[i], pos)
+	source := mustGetColPosition(expr)
+	var replace func(*plan.Expr)
+	replace = func(expr *plan.Expr) {
+		switch e := expr.Expr.(type) {
+		case *plan.Expr_Col:
+			if e.Col.ColPos == source {
+				e.Col.ColPos = pos
+			}
+		case *plan.Expr_F:
+			for _, arg := range e.F.Args {
+				replace(arg)
 			}
 		}
 	}
+	replace(expr)
 }
 
 // mustGetColPosition extracts the column position from an expression.
@@ -461,6 +471,40 @@ func mustGetColPosition(expr *plan.Expr) int32 {
 	return -1
 }
 
+// A single equality value can be substituted only when every partition uses
+// that column. Named scans use column identity; legacy nameless scans must also
+// establish that every reference has the same schema position.
+func matchesPartitionColumns(col *plan.ColRef, position int32, metadata partition.PartitionMetadata) bool {
+	for _, pt := range metadata.Partitions {
+		if !matchesPartitionColumn(col, position, pt.Expr) {
+			return false
+		}
+		if col.Name == "" {
+			var samePosition func(*plan.Expr) bool
+			samePosition = func(expr *plan.Expr) bool {
+				if expr == nil {
+					return false
+				}
+				if ref := expr.GetCol(); ref != nil {
+					return ref.ColPos == position
+				}
+				if fn := expr.GetF(); fn != nil {
+					for _, arg := range fn.Args {
+						if !samePosition(arg) {
+							return false
+						}
+					}
+				}
+				return true
+			}
+			if !samePosition(pt.Expr) {
+				return false
+			}
+		}
+	}
+	return true
+}
+
 // Scan column positions are compacted independently of the stored partition
 // expression. Use column identity when the scan supplies a name; retain the
 // positional fallback for older nameless expressions.

pkg/partitionprune/filter_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/container/vector"
 	"github.com/matrixorigin/matrixone/pkg/pb/partition"
 	"github.com/matrixorigin/matrixone/pkg/pb/plan"
+	p "github.com/matrixorigin/matrixone/pkg/sql/plan"
 	"github.com/matrixorigin/matrixone/pkg/vm/process"
 	"github.com/stretchr/testify/require"
 )
@@ -1711,3 +1712,50 @@ func TestIntersectSortedSlices(t *testing.T) {
 		})
 	}
 }
+
+func TestFilterRepeatedPartitionColumn(t *testing.T) {
+	mp := mpool.MustNewZeroNoFixed()
+	proc := process.NewTopProcess(context.Background(), mp, nil, nil, nil, nil, nil, nil, nil, nil, nil)
+	defer proc.Free()
+	literal := func(v int64) *plan.Expr {
+		return &plan.Expr{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Lit{Lit: &plan.Literal{Value: &plan.Literal_I64Val{I64Val: v}}}}
+	}
+	for _, method := range []partition.PartitionMethod{partition.PartitionMethod_Range, partition.PartitionMethod_Hash, partition.PartitionMethod_Key} {
+		t.Run(method.String(), func(t *testing.T) {
+			sum := bindPartitionTestExpr(t, "+", partitionTestColumn(0), partitionTestColumn(0))
+			nested := bindPartitionTestExpr(t, "*", partitionTestColumn(0), sum)
+			metadata := partition.PartitionMetadata{Method: method, Partitions: []partition.Partition{
+				{Position: 0, Expr: bindPartitionTestExpr(t, "<", nested, literal(0))},
+				{Position: 1, Expr: bindPartitionTestExpr(t, ">=", nested, literal(0))},
+			}}
+			filter := bindPartitionTestExpr(t, "=", partitionTestColumn(0), literal(1))
+			beforeFilter := p.DeepCopyExpr(filter)
+			beforeFirst, beforeSecond := p.DeepCopyExpr(metadata.Partitions[0].Expr), p.DeepCopyExpr(metadata.Partitions[1].Expr)
+			targets, err := Filter(proc, []*plan.Expr{filter}, metadata)
+			require.NoError(t, err)
+			require.Equal(t, []int{1}, targets)
+			require.Equal(t, beforeFilter, filter)
+			require.Equal(t, beforeFirst, metadata.Partitions[0].Expr)
+			require.Equal(t, beforeSecond, metadata.Partitions[1].Expr)
+			// A later partition with another column cannot be substituted by
+			// the same value, including legacy nameless filter expressions.
+			other := partitionTestColumn(1)
+			other.GetCol().Name = "other"
+			metadata.Partitions[1].Expr = bindPartitionTestExpr(t, ">=", other, literal(0))
+			for _, name := range []string{"id", ""} {
+				filter.GetF().Args[0].GetCol().Name = name
+				targets, err = Filter(proc, []*plan.Expr{filter}, metadata)
+				require.NoError(t, err)
+				require.Equal(t, []int{0, 1}, targets)
+			}
+		})
+	}
+	root, err := mustReplaceCol(proc.Ctx, partitionTestColumn(0), literal(1))
+	require.NoError(t, err)
+	require.Equal(t, "cast", root.GetF().Func.ObjName)
+	invalid := partitionTestColumn(0)
+	invalid.Typ.Id = int32(types.T_Rowid)
+	root, err = mustReplaceCol(proc.Ctx, invalid, literal(1))
+	require.Error(t, err)
+	require.Nil(t, root)
+}

pkg/partitionprune/prune_test.go

@@ -16,6 +16,7 @@ package partitionprune
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/matrixorigin/matrixone/pkg/common/mpool"
@@ -26,6 +27,7 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/container/vector"
 	"github.com/matrixorigin/matrixone/pkg/pb/partition"
 	"github.com/matrixorigin/matrixone/pkg/pb/plan"
+	p "github.com/matrixorigin/matrixone/pkg/sql/plan"
 	"github.com/matrixorigin/matrixone/pkg/vm/process"
 	"github.com/stretchr/testify/require"
 )
@@ -364,3 +366,137 @@ func batchCount(result partitionservice.PruneResult) int {
 	})
 	return count
 }
+
+// Bind real overloads so these tests exercise evaluation with incompatible
+// incoming and resolved identities, rather than only inspecting rewritten ASTs.
+func bindPartitionTestExpr(t *testing.T, name string, args ...*plan.Expr) *plan.Expr {
+	t.Helper()
+	expr, err := p.BindFuncExprImplByPlanExpr(context.Background(), name, args)
+	require.NoError(t, err)
+	return expr
+}
+
+func partitionTestColumn(pos int32) *plan.Expr {
+	return &plan.Expr{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: pos, Name: "id"}}}
+}
+
+func TestPruneRepeatedResolvedIdentity(t *testing.T) {
+	for _, shape := range []string{"sum", "square", "nested"} {
+		for _, rowID := range []bool{false, true} {
+			t.Run(fmt.Sprintf("%s/rowid=%v", shape, rowID), func(t *testing.T) {
+				mp := mpool.MustNewZeroNoFixed()
+				proc := process.NewTopProcess(context.Background(), mp, nil, nil, nil, nil, nil, nil, nil, nil, nil)
+				defer proc.Free()
+				build := func(pos int32) *plan.Expr {
+					value := bindPartitionTestExpr(t, "+", partitionTestColumn(pos), partitionTestColumn(pos))
+					if shape == "square" {
+						value = bindPartitionTestExpr(t, "*", partitionTestColumn(pos), partitionTestColumn(pos))
+					}
+					if shape == "nested" {
+						value = bindPartitionTestExpr(t, "*", partitionTestColumn(pos), value)
+					}
+					zero := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Lit{Lit: &plan.Literal{Value: &plan.Literal_I64Val{I64Val: 0}}}}
+					return bindPartitionTestExpr(t, ">=", value, zero)
+
+				}
+				expr, shifted := build(0), build(1)
+				original := p.DeepCopyExpr(expr)
+				originalShifted := p.DeepCopyExpr(shifted)
+				bat := batch.NewWithSize(3)
+				index := 2
+				if rowID {
+					bat = batch.NewWithSize(4)
+					bat.Vecs[0] = vector.NewVec(types.T_Rowid.ToType())
+					require.NoError(t, vector.AppendFixed(bat.Vecs[0], types.Rowid{}, false, mp))
+					index++
+				}
+				defer bat.Clean(mp)
+				offset := 0
+				if rowID {
+					offset = 1
+				}
+				for i, v := range []int64{-1, 42, 1} {
+					bat.Vecs[i+offset] = vector.NewVec(types.T_int64.ToType())
+					require.NoError(t, vector.AppendFixed(bat.Vecs[i+offset], v, false, mp))
+				}
+				bat.SetRowCount(1)
+				// partitionIndex excludes RowID, while index selects the actual
+				// resolved vector in the batch.
+				result, err := PrunePartitionByExpr(proc, bat, partition.Partition{Expr: expr, ExprWithRowID: shifted}, 2)
+				require.NoError(t, err)
+				defer result.Clean(mp)
+				require.Equal(t, 1, result.RowCount())
+				require.Equal(t, []int64{1}, vector.MustFixedColWithTypeCheck[int64](result.Vecs[index]))
+				require.Equal(t, original, expr)
+				require.Equal(t, originalShifted, shifted)
+			})
+		}
+	}
+}
+
+func TestPartitionColumnTraversal(t *testing.T) {
+	for _, reverse := range []bool{false, true} {
+		t.Run(fmt.Sprintf("reverse=%v", reverse), func(t *testing.T) {
+			direct := partitionTestColumn(0)
+			nested := bindPartitionTestExpr(t, "+", partitionTestColumn(0), partitionTestColumn(1))
+			args := []*plan.Expr{direct, nested}
+			if reverse {
+				args[0], args[1] = args[1], args[0]
+			}
+			expr := bindPartitionTestExpr(t, "+", args...)
+			// Select the first schema column, preserving the other column.
+			mustReplaceColPos(expr, 3)
+			require.Equal(t, int32(3), direct.GetCol().ColPos)
+			require.Equal(t, int32(3), nested.GetF().Args[0].GetCol().ColPos)
+			require.Equal(t, int32(1), nested.GetF().Args[1].GetCol().ColPos)
+		})
+	}
+	root := partitionTestColumn(2)
+	mustReplaceColPos(root, 4)
+	require.Equal(t, int32(4), root.GetCol().ColPos)
+
+}
+
+func TestPrunePreservesOtherPartitionColumns(t *testing.T) {
+	mp := mpool.MustNewZeroNoFixed()
+	proc := process.NewTopProcess(context.Background(), mp, nil, nil, nil, nil, nil, nil, nil, nil, nil)
+	defer proc.Free()
+	build := func(offset int32) *plan.Expr {
+		other := partitionTestColumn(offset + 1)
+		other.GetCol().Name = "other"
+		sum := bindPartitionTestExpr(t, "+", partitionTestColumn(offset), other)
+		zero := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Lit{Lit: &plan.Literal{Value: &plan.Literal_I64Val{I64Val: 0}}}}
+		return bindPartitionTestExpr(t, "<", sum, zero)
+
+	}
+	expr, shifted := build(0), build(1)
+	for _, rowID := range []bool{false, true} {
+		for _, explicit := range []int32{-1, 2} {
+			t.Run(fmt.Sprintf("rowid=%v/index=%d", rowID, explicit), func(t *testing.T) {
+				offset := 0
+				if rowID {
+					offset = 1
+				}
+				bat := batch.NewWithSize(3 + offset)
+				defer bat.Clean(mp)
+				if rowID {
+					bat.Vecs[0] = vector.NewVec(types.T_Rowid.ToType())
+					require.NoError(t, vector.AppendFixed(bat.Vecs[0], types.Rowid{}, false, mp))
+				}
+				for i, v := range []int64{-1, -100, 1} {
+					bat.Vecs[i+offset] = vector.NewVec(types.T_int64.ToType())
+					require.NoError(t, vector.AppendFixed(bat.Vecs[i+offset], v, false, mp))
+				}
+				bat.SetRowCount(1)
+				before, beforeShifted := p.DeepCopyExpr(expr), p.DeepCopyExpr(shifted)
+				result, err := PrunePartitionByExpr(proc, bat, partition.Partition{Expr: expr, ExprWithRowID: shifted}, explicit)
+				require.NoError(t, err)
+				defer result.Clean(mp)
+				require.Equal(t, 1, result.RowCount(), "unrelated column must remain -100")
+				require.Equal(t, []int64{-100}, vector.MustFixedColWithTypeCheck[int64](result.Vecs[1+offset]))
+				require.Equal(t, before, expr)
+				require.Equal(t, beforeShifted, shifted)
+			})
+		}
+	}
+}

pkg/partitionservice/storage.go

@@ -1064,15 +1064,11 @@ func getExprWithRowID(v []byte) *plan.Expr {
 
 func resetPosWithRowID(expr *plan.Expr) {
 	switch e := expr.Expr.(type) {
+	case *plan.Expr_Col:
+		e.Col.ColPos++
 	case *plan.Expr_F:
-		for i := range e.F.Args {
-			switch col := e.F.Args[i].Expr.(type) {
-			case *plan.Expr_Col:
-				col.Col.ColPos++
-				return
-			case *plan.Expr_F:
-				resetPosWithRowID(e.F.Args[i])
-			}
+		for _, arg := range e.F.Args {
+			resetPosWithRowID(arg)
 		}
 	}
 }