perf(readutil): preserve early filtering across fallback paths

pkg/vm/engine/disttae/local_disttae_datasource.go

@@ -775,10 +775,9 @@ func (ls *LocalDisttaeDataSource) filterInMemCommittedInserts(
 		if !ls.memPKFilter.Valid() {
 			ls.pStateRows.insIter = ls.pState.NewRowsIter(ls.snapshotTS, nil, false)
 		} else {
-			ls.pStateRows.insIter = ls.pState.NewPrimaryKeyIter(
+			ls.pStateRows.insIter = ls.pState.NewPrimaryKeyIterWithFilters(
 				ls.memPKFilter.TS,
-				ls.memPKFilter.Op(),
-				ls.memPKFilter.Keys())
+				ls.memPKFilter.Specs())
 		}
 		if summaryBuf != nil {
 			summaryBuf.WriteString(fmt.Sprintf("[PScan] insIter created %v\n", ls.memPKFilter.String()))
@@ -1287,6 +1286,14 @@ func (ls *LocalDisttaeDataSource) getInMemDelIter(
 		ls.memPKFilter == nil || !ls.memPKFilter.Valid() {
 		return ls.pState.NewRowsIter(ls.snapshotTS, bid, true), false
 	}
+	filterSpecs := ls.memPKFilter.Specs()
+	if len(filterSpecs) > 1 {
+		return ls.pState.NewPrimaryKeyDelIterWithFilters(
+			&ls.memPKFilter.TS,
+			bid,
+			filterSpecs,
+		), false
+	}
 
 	inValCnt, ok := ls.memPKFilter.InKind()
 	if !ok {

pkg/vm/engine/disttae/logtailreplay/rows_iter.go

@@ -16,6 +16,7 @@ package logtailreplay
 
 import (
 	"bytes"
+	"container/heap"
 	"fmt"
 	"math"
 
@@ -814,6 +815,151 @@ func (p *PartitionState) NewPrimaryKeyIter(
 	}
 }
 
+type primaryKeyIterHeap struct {
+	children []primaryKeyOrderedIter
+	indexes  []int
+}
+
+func (h primaryKeyIterHeap) Len() int {
+	return len(h.indexes)
+}
+
+func (h primaryKeyIterHeap) Less(i, j int) bool {
+	left := h.children[h.indexes[i]].currentIndexEntry()
+	right := h.children[h.indexes[j]].currentIndexEntry()
+	return left.Less(right)
+}
+
+func (h primaryKeyIterHeap) Swap(i, j int) {
+	h.indexes[i], h.indexes[j] = h.indexes[j], h.indexes[i]
+}
+
+func (h *primaryKeyIterHeap) Push(value any) {
+	h.indexes = append(h.indexes, value.(int))
+}
+
+func (h *primaryKeyIterHeap) Pop() any {
+	last := len(h.indexes) - 1
+	value := h.indexes[last]
+	h.indexes = h.indexes[:last]
+	return value
+}
+
+type primaryKeyUnionIter struct {
+	heap        primaryKeyIterHeap
+	initialized bool
+	closed      bool
+	curRow      *RowEntry
+	last        PrimaryIndexEntry
+	hasLast     bool
+}
+
+var _ RowsIter = new(primaryKeyUnionIter)
+
+type primaryKeyOrderedIter interface {
+	RowsIter
+	currentIndexEntry() *PrimaryIndexEntry
+}
+
+func (p *primaryKeyIter) currentIndexEntry() *PrimaryIndexEntry {
+	return p.iter.Item()
+}
+
+func samePrimaryIndexEntry(left, right *PrimaryIndexEntry) bool {
+	return left.RowEntryID == right.RowEntryID &&
+		left.Time.Compare(&right.Time) == 0 &&
+		bytes.Equal(left.Bytes, right.Bytes)
+}
+
+func (p *primaryKeyUnionIter) Next() bool {
+	if p.closed {
+		return false
+	}
+	if !p.initialized {
+		p.initialized = true
+		for idx := range p.heap.children {
+			if p.heap.children[idx].Next() {
+				heap.Push(&p.heap, idx)
+			}
+		}
+	}
+
+	for p.heap.Len() > 0 {
+		idx := heap.Pop(&p.heap).(int)
+		child := p.heap.children[idx]
+		entry := *child.currentIndexEntry()
+		row := child.Entry()
+		if child.Next() {
+			heap.Push(&p.heap, idx)
+		}
+
+		if p.hasLast && samePrimaryIndexEntry(&entry, &p.last) {
+			continue
+		}
+		p.last = entry
+		p.hasLast = true
+		p.curRow = row
+		return true
+	}
+
+	return false
+}
+
+func (p *primaryKeyUnionIter) Entry() *RowEntry {
+	return p.curRow
+}
+
+func (p *primaryKeyUnionIter) Close() error {
+	if p.closed {
+		return nil
+	}
+	p.closed = true
+	var firstErr error
+	for idx := range p.heap.children {
+		if err := p.heap.children[idx].Close(); err != nil && firstErr == nil {
+			firstErr = err
+		}
+	}
+	p.heap.children = nil
+	p.heap.indexes = nil
+	return firstErr
+}
+
+func (p *PartitionState) NewPrimaryKeyIterWithFilters(
+	ts types.TS,
+	filters []readutil.MemPKFilterSpec,
+) RowsIter {
+	if len(filters) == 1 {
+		return p.NewPrimaryKeyIter(ts, filters[0].Op, filters[0].Keys)
+	}
+
+	children := make([]primaryKeyOrderedIter, 0, len(filters))
+	for idx := range filters {
+		children = append(children, p.NewPrimaryKeyIter(ts, filters[idx].Op, filters[idx].Keys))
+	}
+	return &primaryKeyUnionIter{
+		heap: primaryKeyIterHeap{children: children},
+	}
+}
+
+func (p *PartitionState) NewPrimaryKeyDelIterWithFilters(
+	ts *types.TS,
+	bid *types.Blockid,
+	filters []readutil.MemPKFilterSpec,
+) RowsIter {
+	if len(filters) == 1 {
+		return p.NewPrimaryKeyDelIter(ts, bid, filters[0].Op, filters[0].Keys)
+	}
+
+	children := make([]primaryKeyOrderedIter, 0, len(filters))
+	for idx := range filters {
+		children = append(children, p.NewPrimaryKeyDelIter(ts, bid, filters[idx].Op, filters[idx].Keys))
+	}
+	return &primaryKeyUnionIter{
+		heap: primaryKeyIterHeap{children: children},
+	}
+}
+
 func (p *PartitionState) NewPrimaryKeyDelIter(
 	ts *types.TS,
 	bid *types.Blockid,

pkg/vm/engine/disttae/logtailreplay/rows_iter_test.go

@@ -569,6 +569,94 @@ func TestPrefixIn(t *testing.T) {
 	require.Equal(t, []byte{4}, pkIter.iter.Item().Bytes)
 }
 
+func TestPrimaryKeyUnionIter(t *testing.T) {
+	state := NewPartitionState("", false, 42, false)
+	pool := mpool.MustNewZero()
+	packer := types.NewPacker()
+	defer packer.Close()
+
+	rowIDVec := vector.NewVec(types.T_Rowid.ToType())
+	tsVec := vector.NewVec(types.T_TS.ToType())
+	pkVec := vector.NewVec(types.T_int64.ToType())
+	segmentID := objectio.NewSegmentid()
+	blockIDs := make([]*types.Blockid, 0, 6)
+	for idx := 0; idx < 6; idx++ {
+		blockID := objectio.NewBlockid(segmentID, uint16(idx), 0)
+		blockIDs = append(blockIDs, blockID)
+		vector.AppendFixed(rowIDVec, objectio.NewRowid(blockID, 0), false, pool)
+		vector.AppendFixed(tsVec, types.BuildTS(1, 0), false, pool)
+		vector.AppendFixed(pkVec, int64(idx), false, pool)
+	}
+	state.HandleRowsInsert(context.Background(), &api.Batch{
+		Attrs: []string{"rowid", "time", "pk"},
+		Vecs: []api.Vector{
+			mustVectorToProto(rowIDVec),
+			mustVectorToProto(tsVec),
+			mustVectorToProto(pkVec),
+		},
+	}, 0, packer, pool)
+
+	iter := state.NewPrimaryKeyIterWithFilters(types.BuildTS(2, 0), []readutil.MemPKFilterSpec{
+		{Op: function.EQUAL, Keys: [][]byte{readutil.EncodePrimaryKey(int64(1), packer)}},
+		{Op: function.BETWEEN, Keys: [][]byte{
+			readutil.EncodePrimaryKey(int64(1), packer),
+			readutil.EncodePrimaryKey(int64(3), packer),
+		}},
+	})
+	union, ok := iter.(*primaryKeyUnionIter)
+	require.True(t, ok)
+	require.Len(t, union.heap.children, 2)
+
+	var got []int64
+	for iter.Next() {
+		values, err := types.Unpack(iter.Entry().PrimaryIndexBytes)
+		require.NoError(t, err)
+		got = append(got, values[0].(int64))
+	}
+	require.Equal(t, []int64{1, 2, 3}, got)
+	require.NoError(t, iter.Close())
+	require.NoError(t, iter.Close())
+	require.False(t, iter.Next())
+
+	empty := state.NewPrimaryKeyIterWithFilters(types.BuildTS(2, 0), nil)
+	require.False(t, empty.Next())
+	require.NoError(t, empty.Close())
+
+	deletedRowIDs := vector.NewVec(types.T_Rowid.ToType())
+	deleteTS := vector.NewVec(types.T_TS.ToType())
+	deletedPKs := vector.NewVec(types.T_int64.ToType())
+	tombstoneRowIDs := vector.NewVec(types.T_Rowid.ToType())
+	vector.AppendFixed(deletedRowIDs, objectio.NewRowid(blockIDs[1], 0), false, pool)
+	vector.AppendFixed(deleteTS, types.BuildTS(3, 0), false, pool)
+	vector.AppendFixed(deletedPKs, int64(1), false, pool)
+	vector.AppendFixed(tombstoneRowIDs, types.RandomRowid(), false, pool)
+	state.HandleRowsDelete(context.Background(), &api.Batch{
+		Attrs: []string{"rowid", "time"},
+		Vecs: []api.Vector{
+			mustVectorToProto(deletedRowIDs),
+			mustVectorToProto(deleteTS),
+			mustVectorToProto(deletedPKs),
+			mustVectorToProto(tombstoneRowIDs),
+		},
+	}, packer, pool)
+
+	deleteSnapshot := types.BuildTS(4, 0)
+	deleteIter := state.NewPrimaryKeyDelIterWithFilters(
+		&deleteSnapshot,
+		blockIDs[1],
+		[]readutil.MemPKFilterSpec{
+			{Op: function.EQUAL, Keys: [][]byte{readutil.EncodePrimaryKey(int64(1), packer)}},
+			{Op: function.BETWEEN, Keys: [][]byte{
+				readutil.EncodePrimaryKey(int64(1), packer),
+				readutil.EncodePrimaryKey(int64(3), packer),
+			}},
+		},
+	)
+	require.True(t, deleteIter.Next())
+	require.False(t, deleteIter.Next())
+	require.NoError(t, deleteIter.Close())
+}
+
 func TestPrefixBetweenOpenBounds(t *testing.T) {
 	pkTree := btree.NewBTreeGOptions((*PrimaryIndexEntry).Less, btree.Options{Degree: 64})
 	for _, key := range []string{"aa1", "aa2", "ab1", "ab2", "ac1"} {

pkg/vm/engine/disttae/pk_filter_observe.go

@@ -43,13 +43,21 @@ func logPKFilterInMemSummary(
 		filterValid bool
 		filterOp    int
 		keyCnt      int
+		disjunctCnt int
 		exact       bool
 		exactHit    bool
 	)
 	if filter != nil {
 		filterValid = filter.Valid()
 		filterOp = filter.Op()
-		keyCnt = len(filter.Keys())
+		specs := filter.Specs()
+		disjunctCnt = len(specs)
+		for idx := range specs {
+			keyCnt += len(specs[idx].Keys)
+		}
+		if disjunctCnt > 1 {
+			filterOp = -1
+		}
 		exact, exactHit = filter.Exact()
 	}
 
@@ -61,6 +69,7 @@ func logPKFilterInMemSummary(
 		zap.Bool("has-mem-pk-filter", filter != nil),
 		zap.Bool("mem-pk-filter-valid", filterValid),
 		zap.Int("filter-op", filterOp),
+		zap.Int("disjunct-count", disjunctCnt),
 		zap.Int("key-count", keyCnt),
 		zap.Bool("exact-filter", exact),
 		zap.Bool("exact-hit", exactHit),

pkg/vm/engine/disttae/snapshot_materialized_datasource.go

@@ -214,10 +214,9 @@ func (ds *materializedSnapshotDataSource) buildCommittedInMemBatches(
 	iterKind := "rows"
 	if ds.memPKFilter != nil && ds.memPKFilter.Valid() {
 		iterKind = "primary-key"
-		iter = ds.pState.NewPrimaryKeyIter(
+		iter = ds.pState.NewPrimaryKeyIterWithFilters(
 			ds.snapshotTS,
-			ds.memPKFilter.Op(),
-			ds.memPKFilter.Keys(),
+			ds.memPKFilter.Specs(),
 		)
 	} else {
 		iter = ds.pState.NewRowsIter(ds.snapshotTS, nil, false)

pkg/vm/engine/readutil/expr_filter.go

@@ -125,12 +125,14 @@ func CompileFilterExprs(
 	ops3 := make([]ObjectFilterOp, 0, len(exprs))
 	ops4 := make([]BlockFilterOp, 0, len(exprs))
 	ops5 := make([]SeekFirstBlockOp, 0, len(exprs))
+	compiled := 0
 
 	for _, expr := range exprs {
 		expr_op1, expr_op2, expr_op3, expr_op4, expr_op5, can, hsh := CompileFilterExpr(expr, tableDef, fs)
 		if !can {
-			return nil, nil, nil, nil, nil, false, false
+			continue
 		}
+		compiled++
 		if expr_op1 != nil {
 			ops1 = append(ops1, expr_op1)
 		}
@@ -148,6 +150,9 @@ func CompileFilterExprs(
 		}
 		highSelectivityHint = highSelectivityHint || hsh
 	}
+	if compiled == 0 {
+		return nil, nil, nil, nil, nil, false, false
+	}
 	fastFilterOp = func(obj *objectio.ObjectStats) (bool, error) {
 		for _, op := range ops1 {
 			ok, err := op(obj)
@@ -335,12 +340,14 @@ func CompileFilterExpr(
 			objectOps := make([]ObjectFilterOp, 0, len(exprImpl.F.Args))
 			blockOps := make([]BlockFilterOp, 0, len(exprImpl.F.Args))
 			seekOps := make([]SeekFirstBlockOp, 0, len(exprImpl.F.Args))
+			compiled := 0
 
 			for idx := range exprImpl.F.Args {
 				op1, op2, op3, op4, op5, can, hsh := CompileFilterExpr(exprImpl.F.Args[idx], tableDef, fs)
 				if !can {
-					return nil, nil, nil, nil, nil, false, false
+					continue
 				}
+				compiled++
 
 				fastOps = append(fastOps, op1)
 				loadOps = append(loadOps, op2)
@@ -350,6 +357,9 @@ func CompileFilterExpr(
 
 				highSelectivityHint = highSelectivityHint || hsh
 			}
+			if compiled == 0 {
+				return nil, nil, nil, nil, nil, false, false
+			}
 
 			fastFilterOp = func(stats *objectio.ObjectStats) (bool, error) {
 				for idx := range fastOps {