fix: column pruning failure when applying associative law rules (#23206)

This PR fixes a column pruning failure issue caused by applying associative law optimization rules. When the query optimizer applies associative law transformations to reorder JOIN operations, the OnList conditions were not properly migrated to match the new JOIN structure, leading to remapping errors during column reference resolution.

### Problem

When applying associative law rules (Rule 1, Rule 2, Rule 3) to transform JOIN structures like `A*(B*C)` to `(A*B)*C`, the OnList conditions that reference columns from moved tables were not migrated to the correct JOIN nodes. This caused column remapping failures because:

1. The JOIN structure was changed, but OnList conditions remained in their original positions
2. During column remapping phase, references to columns from moved tables could not be found in the expected context
3. This resulted in errors like "Column remapping failed: cannot find column reference"

### Solution

The fix adds proper OnList condition migration logic in all three associative law rules:

1. **Added `migrateOnListConditions` function**: Moves conditions involving specified table tags from source node to destination node
2. **Added `checkExprInvolvesTags` function**: Checks if an expression references any of the specified table tags
3. **Updated Rule 1, Rule 2, Rule 3**: Each rule now properly migrates OnList conditions when transforming JOIN structures:
- **Rule 1**: Migrates conditions involving table C to the outer join
- **Rule 2**: Migrates conditions involving table A to the outer join
- **Rule 3**: Migrates conditions involving B and C to their correct positions

### Testing

Added a comprehensive test case `TestAssociativeLawRemapping` that:
- Creates a scenario where associative law Rule 1 would be triggered
- Verifies that the query executes successfully without remapping errors
- Tests the exact scenario from the issue report

Approved by: @aunjgr, @XuPeng-SH

Author: aptend

pkg/sql/compile/sql_executor.go

@@ -190,13 +190,18 @@ func (s *sqlExecutor) getCompileContext(
 	proc *process.Process,
 	db string,
 	lower int64) *compilerContext {
-	return &compilerContext{
+	cc := &compilerContext{
 		ctx:       ctx,
 		defaultDB: db,
 		engine:    s.eng,
 		proc:      proc,
 		lower:     lower,
 	}
+	// For testing: check if a stats cache is provided in context
+	if statsCache, ok := ctx.Value("test_stats_cache").(*plan.StatsCache); ok {
+		cc.statsCache = statsCache
+	}
+	return cc
 }
 
 func (s *sqlExecutor) adjustOptions(

pkg/sql/compile/sql_executor_context.go

@@ -24,6 +24,7 @@ import (
 
 	"github.com/matrixorigin/matrixone/pkg/common/moerr"
 	"github.com/matrixorigin/matrixone/pkg/defines"
+	"github.com/matrixorigin/matrixone/pkg/logutil"
 	planpb "github.com/matrixorigin/matrixone/pkg/pb/plan"
 	pb "github.com/matrixorigin/matrixone/pkg/pb/statsinfo"
 	"github.com/matrixorigin/matrixone/pkg/perfcounter"
@@ -137,6 +138,17 @@ func (c *compilerContext) Stats(obj *plan.ObjectRef, snapshot *plan.Snapshot) (*
 		return nil, moerr.NewNoSuchTable(ctx, dbName, tableName)
 	}
 
+	// For testing only: check if stats cache already has stats for this table
+	// Only use cached stats if we're in a test environment (indicated by test_stats_cache in context)
+	// If so, use the cached stats instead of fetching from engine
+	if _, isTestEnv := c.GetContext().Value("test_stats_cache").(*plan.StatsCache); isTestEnv {
+		tableID := table.GetTableID(ctx)
+		if statsWrapper := c.statsCache.GetStatsInfo(tableID, false); statsWrapper != nil && statsWrapper.Stats != nil {
+			logutil.Infof("use test env cached stats for table %s (tableID=%d)", tableName, tableID)
+			return statsWrapper.Stats, nil
+		}
+	}
+
 	newCtx := perfcounter.AttachCalcTableStatsKey(ctx)
 	statsInfo, err := table.Stats(newCtx, true)
 	if err != nil {

pkg/sql/plan/associative_law.go

@@ -14,7 +14,65 @@
 
 package plan
 
-import "github.com/matrixorigin/matrixone/pkg/pb/plan"
+import (
+	"fmt"
+
+	"github.com/matrixorigin/matrixone/pkg/pb/plan"
+)
+
+// checkExprInvolvesTags checks if an expression references any of the specified tags
+func checkExprInvolvesTags(expr *plan.Expr, tagsMap map[int32]bool) bool {
+	switch e := expr.Expr.(type) {
+	case *plan.Expr_Col:
+		return tagsMap[e.Col.RelPos]
+	case *plan.Expr_F:
+		for _, arg := range e.F.Args {
+			if checkExprInvolvesTags(arg, tagsMap) {
+				return true
+			}
+		}
+	case *plan.Expr_W:
+		if checkExprInvolvesTags(e.W.WindowFunc, tagsMap) {
+			return true
+		}
+		for _, order := range e.W.OrderBy {
+			if checkExprInvolvesTags(order.Expr, tagsMap) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// migrateOnListConditions moves conditions involving specified tags from src to dst
+func migrateOnListConditions(src *plan.Node, dst *plan.Node, tagsMap map[int32]bool) {
+	var kept []*plan.Expr
+	for _, cond := range src.OnList {
+		if checkExprInvolvesTags(cond, tagsMap) {
+			dst.OnList = append(dst.OnList, cond)
+		} else {
+			kept = append(kept, cond)
+		}
+	}
+	src.OnList = kept
+}
+
+// getTableNameOrLabel returns table name from node's TableDef, or a label (A/B/C) if not available
+func (builder *QueryBuilder) getTableNameOrLabel(nodeID int32, label string) string {
+	node := builder.qry.Nodes[nodeID]
+	if node.TableDef != nil && node.TableDef.Name != "" {
+		return node.TableDef.Name
+	}
+	return label
+}
+
+// formatStatsInfo formats stats information for logging
+func formatStatsInfo(stats *plan.Stats) string {
+	if stats == nil {
+		return "stats=nil"
+	}
+	return fmt.Sprintf("sel=%.4f,outcnt=%.2f,tablecnt=%.2f", stats.Selectivity, stats.Outcnt, stats.TableCnt)
+}
 
 // for A*(B*C), if C.sel>0.9 and B<C, change this to (A*B)*C
 func (builder *QueryBuilder) applyAssociativeLawRule1(nodeID int32) int32 {
@@ -43,6 +101,25 @@ func (builder *QueryBuilder) applyAssociativeLawRule1(nodeID int32) int32 {
 		node.Children[1] = rightChild.NodeId
 		return node.NodeId
 	}
+
+	tagsC := builder.enumerateTags(NodeC.NodeId)
+	// Migrate OnList: conditions involving C must move to outer join
+	tagsCMap := make(map[int32]bool)
+	for _, tag := range tagsC {
+		tagsCMap[tag] = true
+	}
+	migrateOnListConditions(node, rightChild, tagsCMap)
+
+	// Record table names and stats after migration
+	tableNameA := builder.getTableNameOrLabel(node.Children[0], "A")
+	tableNameB := builder.getTableNameOrLabel(NodeB.NodeId, "B")
+	tableNameC := builder.getTableNameOrLabel(NodeC.NodeId, "C")
+	statsInfo := fmt.Sprintf("rule1: A=%s(stats:%s) B=%s(stats:%s) C=%s(stats:%s)",
+		tableNameA, formatStatsInfo(builder.qry.Nodes[node.Children[0]].Stats),
+		tableNameB, formatStatsInfo(NodeB.Stats),
+		tableNameC, formatStatsInfo(NodeC.Stats))
+	builder.optimizationHistory = append(builder.optimizationHistory, statsInfo)
+
 	rightChild.Children[0] = node.NodeId
 	ReCalcNodeStats(rightChild.NodeId, builder, true, false, true)
 	return rightChild.NodeId
@@ -67,6 +144,7 @@ func (builder *QueryBuilder) applyAssociativeLawRule2(nodeID int32) int32 {
 	if NodeC.Stats.Selectivity > 0.5 {
 		return nodeID
 	}
+	NodeA := builder.qry.Nodes[leftChild.Children[0]]
 	NodeB := builder.qry.Nodes[leftChild.Children[1]]
 	node.Children[0] = NodeB.NodeId
 	determineHashOnPK(node.NodeId, builder)
@@ -75,6 +153,24 @@ func (builder *QueryBuilder) applyAssociativeLawRule2(nodeID int32) int32 {
 		node.Children[0] = leftChild.NodeId
 		return node.NodeId
 	}
+
+	// Migrate OnList: conditions involving A must move to outer join
+	tagsAMap := make(map[int32]bool)
+	for _, tag := range builder.enumerateTags(NodeA.NodeId) {
+		tagsAMap[tag] = true
+	}
+	migrateOnListConditions(node, leftChild, tagsAMap)
+
+	// Record table names and stats after migration
+	tableNameA := builder.getTableNameOrLabel(NodeA.NodeId, "A")
+	tableNameB := builder.getTableNameOrLabel(NodeB.NodeId, "B")
+	tableNameC := builder.getTableNameOrLabel(NodeC.NodeId, "C")
+	statsInfo := fmt.Sprintf("rule2: A=%s(stats:%s) B=%s(stats:%s) C=%s(stats:%s)",
+		tableNameA, formatStatsInfo(NodeA.Stats),
+		tableNameB, formatStatsInfo(NodeB.Stats),
+		tableNameC, formatStatsInfo(NodeC.Stats))
+	builder.optimizationHistory = append(builder.optimizationHistory, statsInfo)
+
 	leftChild.Children[1] = node.NodeId
 	ReCalcNodeStats(leftChild.NodeId, builder, true, false, true)
 	return leftChild.NodeId
@@ -110,6 +206,31 @@ func (builder *QueryBuilder) applyAssociativeLawRule3(nodeID int32) int32 {
 		node.Children[0] = leftChild.NodeId
 		return node.NodeId
 	}
+
+	// Migrate OnList:
+	// - node: move conditions involving B to leftChild
+	// - leftChild: move conditions involving C to node
+	tagsBMap := make(map[int32]bool)
+	for _, tag := range builder.enumerateTags(NodeB.NodeId) {
+		tagsBMap[tag] = true
+	}
+	tagsCMap := make(map[int32]bool)
+	for _, tag := range builder.enumerateTags(NodeC.NodeId) {
+		tagsCMap[tag] = true
+	}
+	migrateOnListConditions(node, leftChild, tagsBMap)
+	migrateOnListConditions(leftChild, node, tagsCMap)
+
+	// Record table names and stats after migration
+	tableNameA := builder.getTableNameOrLabel(NodeA.NodeId, "A")
+	tableNameB := builder.getTableNameOrLabel(NodeB.NodeId, "B")
+	tableNameC := builder.getTableNameOrLabel(NodeC.NodeId, "C")
+	statsInfo := fmt.Sprintf("rule3: A=%s(stats:%s) B=%s(stats:%s) C=%s(stats:%s)",
+		tableNameA, formatStatsInfo(NodeA.Stats),
+		tableNameB, formatStatsInfo(NodeB.Stats),
+		tableNameC, formatStatsInfo(NodeC.Stats))
+	builder.optimizationHistory = append(builder.optimizationHistory, statsInfo)
+
 	leftChild.Children[0] = node.NodeId
 	ReCalcNodeStats(leftChild.NodeId, builder, true, false, true)
 	return leftChild.NodeId

pkg/sql/plan/join_order.go

@@ -15,6 +15,7 @@
 package plan
 
 import (
+	"fmt"
 	"sort"
 	"strings"
 
@@ -204,6 +205,7 @@ func HasColExpr(expr *plan.Expr, pos int32) int32 {
 }
 
 func (builder *QueryBuilder) determineJoinOrder(nodeID int32) int32 {
+	originalNodeID := nodeID
 	if builder.optimizerHints != nil && builder.optimizerHints.joinOrdering != 0 {
 		return nodeID
 	}
@@ -223,6 +225,9 @@ func (builder *QueryBuilder) determineJoinOrder(nodeID int32) int32 {
 	}
 
 	leaves, conds := builder.gatherJoinLeavesAndConds(node, nil, nil)
+	// Record middle: gathered leaves and conditions
+	builder.optimizationHistory = append(builder.optimizationHistory,
+		fmt.Sprintf("determineJoinOrder:middle (nodeID: %d, leaves: %d, conds: %d)", nodeID, len(leaves), len(conds)))
 	newConds := deduceNewOnList(conds)
 	conds = append(conds, newConds...)
 	vertices := builder.getJoinGraph(leaves, conds)
@@ -347,6 +352,14 @@ func (builder *QueryBuilder) determineJoinOrder(nodeID int32) int32 {
 			FilterList: conds,
 		}, nil)
 	}
+	// Record after determineJoinOrder
+	if nodeID != originalNodeID {
+		builder.optimizationHistory = append(builder.optimizationHistory,
+			fmt.Sprintf("determineJoinOrder:after (nodeID: %d -> %d, remainingConds: %d)", originalNodeID, nodeID, len(conds)))
+	} else {
+		builder.optimizationHistory = append(builder.optimizationHistory,
+			fmt.Sprintf("determineJoinOrder:after (nodeID: %d, no change, remainingConds: %d)", nodeID, len(conds)))
+	}
 	return nodeID
 }
 

pkg/sql/plan/opt_misc_test.go

@@ -15,10 +15,12 @@
 package plan
 
 import (
+	"context"
+	"testing"
+
 	"github.com/matrixorigin/matrixone/pkg/container/types"
 	"github.com/matrixorigin/matrixone/pkg/pb/plan"
 	"github.com/stretchr/testify/require"
-	"testing"
 )
 
 func TestRemapWindowClause(t *testing.T) {
@@ -43,7 +45,13 @@ func TestRemapWindowClause(t *testing.T) {
 		Typ: plan.Type{},
 	}
 	colMap := make(map[[2]int32][2]int32)
-	var b *QueryBuilder
+	var b *QueryBuilder = &QueryBuilder{
+		compCtx: &MockCompilerContext{
+			ctx: context.Background(),
+		},
+		optimizationHistory: []string{"test optimization history"},
+	}
 	err := b.remapWindowClause(f, 1, 1, colMap, nil)
+	t.Log(err)
 	require.Error(t, err)
 }

pkg/sql/plan/pushdown.go

@@ -15,13 +15,19 @@
 package plan
 
 import (
+	"fmt"
+
 	"github.com/matrixorigin/matrixone/pkg/catalog"
 	"github.com/matrixorigin/matrixone/pkg/objectio"
 	"github.com/matrixorigin/matrixone/pkg/pb/plan"
 	"github.com/matrixorigin/matrixone/pkg/vectorindex/metric"
 )
 
 func (builder *QueryBuilder) pushdownFilters(nodeID int32, filters []*plan.Expr, separateNonEquiConds bool) (int32, []*plan.Expr) {
+	originalNodeID := nodeID
+	// Record before pushdownFilters
+	builder.optimizationHistory = append(builder.optimizationHistory,
+		fmt.Sprintf("pushdownFilters:before (nodeID: %d, nodeType: %s, filters: %d)", nodeID, builder.qry.Nodes[nodeID].NodeType, len(filters)))
 	node := builder.qry.Nodes[nodeID]
 
 	var canPushdown, cantPushdown []*plan.Expr
@@ -150,6 +156,9 @@ func (builder *QueryBuilder) pushdownFilters(nodeID int32, filters []*plan.Expr,
 		}
 
 	case plan.Node_JOIN:
+		// Record middle: processing JOIN node
+		builder.optimizationHistory = append(builder.optimizationHistory,
+			fmt.Sprintf("pushdownFilters:middle (nodeID: %d, JOIN, filters: %d, onList: %d)", nodeID, len(filters), len(node.OnList)))
 		leftTags := make(map[int32]bool)
 		for _, tag := range builder.enumerateTags(node.Children[0]) {
 			leftTags[tag] = true
@@ -402,6 +411,9 @@ func (builder *QueryBuilder) pushdownFilters(nodeID int32, filters []*plan.Expr,
 		node.Children[1] = childID
 
 	case plan.Node_UNION, plan.Node_UNION_ALL, plan.Node_MINUS, plan.Node_MINUS_ALL, plan.Node_INTERSECT, plan.Node_INTERSECT_ALL:
+		// Record middle: processing UNION/MINUS/INTERSECT node
+		builder.optimizationHistory = append(builder.optimizationHistory,
+			fmt.Sprintf("pushdownFilters:middle (nodeID: %d, %s, filters: %d)", nodeID, node.NodeType, len(filters)))
 		leftChild := builder.qry.Nodes[node.Children[0]]
 		rightChild := builder.qry.Nodes[node.Children[1]]
 		var canPushDownRight []*plan.Expr
@@ -457,6 +469,9 @@ func (builder *QueryBuilder) pushdownFilters(nodeID int32, filters []*plan.Expr,
 		node.Children[0] = childID
 
 	case plan.Node_TABLE_SCAN, plan.Node_EXTERNAL_SCAN:
+		// Record middle: processing TABLE_SCAN/EXTERNAL_SCAN node
+		builder.optimizationHistory = append(builder.optimizationHistory,
+			fmt.Sprintf("pushdownFilters:middle (nodeID: %d, %s, filters: %d)", nodeID, node.NodeType, len(filters)))
 		for _, filter := range filters {
 			if onlyContainsTag(filter, node.BindingTags[0]) {
 				node.FilterList = append(node.FilterList, filter)
@@ -505,6 +520,14 @@ func (builder *QueryBuilder) pushdownFilters(nodeID int32, filters []*plan.Expr,
 		}
 	}
 
+	// Record after pushdownFilters
+	if nodeID != originalNodeID {
+		builder.optimizationHistory = append(builder.optimizationHistory,
+			fmt.Sprintf("pushdownFilters:after (nodeID: %d -> %d, cantPushdown: %d)", originalNodeID, nodeID, len(cantPushdown)))
+	} else {
+		builder.optimizationHistory = append(builder.optimizationHistory,
+			fmt.Sprintf("pushdownFilters:after (nodeID: %d, no change, cantPushdown: %d)", nodeID, len(cantPushdown)))
+	}
 	return nodeID, cantPushdown
 }
 

pkg/sql/plan/query_builder.go

@@ -94,6 +94,7 @@ func NewQueryBuilder(queryType plan.Query_StatementType, ctx CompilerContext, is
 		isPrepareStatement:   isPrepareStatement,
 		deleteNode:           make(map[uint64]int32),
 		skipStats:            skipStats,
+		optimizationHistory:  make([]string, 0),
 	}
 }
 
@@ -193,6 +194,17 @@ func (builder *QueryBuilder) buildRemapErrorMessage(
 		}
 	}
 
+	// Optimization history
+	if len(builder.optimizationHistory) > 0 {
+		sb.WriteString("🔧 Optimization History:\n")
+		for _, hist := range builder.optimizationHistory {
+			sb.WriteString(fmt.Sprintf("   - %s\n", hist))
+		}
+		sb.WriteString("\n")
+	} else {
+		sb.WriteString("🔧 Optimization History: (no associatelaw applied)\n\n")
+	}
+
 	// Available columns
 	if len(colMap) > 0 {
 		sb.WriteString("✅ Available Columns in Context:\n")

pkg/sql/plan/types.go

@@ -193,6 +193,10 @@ type QueryBuilder struct {
 	aggSpillMem int64
 
 	optimizerHints *OptimizerHints
+
+	// optimizationHistory records key optimization steps for debugging remap errors
+	// Only records when optimizations actually change the plan structure
+	optimizationHistory []string
 }
 
 type OptimizerHints struct {