[Feature]Sharding support 1

[triump2020]

User description

What type of PR is this?

• API-change
• BUG
• Improvement
• Documentation
• [ x] Feature
• Test and CI
• Code Refactoring

Which issue(s) this PR fixes:

issue #16471

What this PR does / why we need it:

1. Finish collecting tombstones from remote CN.
2. Handle sharding reader

Todo:

• Handle order by
• Improve the performance for sharding.

---

PR Type

Enhancement

---

Description

• Enhanced tombstone collection logic to support both committed and uncommitted tombstones using a policy-based approach.
• Implemented sharding support for tombstone collection, allowing for merging of local and remote tombstones.
• Added new handlers for sharding read operations, including ReadNext, ReadClose, and ReadCollectTombstones.
• Updated mock methods and proto definitions to align with the new sharding and tombstone collection logic.

---

Changes walkthrough 📝

	Relevant files
Enhancement	8 files engine_mock.go Update mock methods with new parameters for tombstone handling pkg/frontend/test/engine_mock.go Updated method signatures for ApplyPersistedTombstones and ApplyTombstones to include additional parameters. Modified mock recorder methods to reflect changes in method signatures. +13/-13  txn_table.go Enhance tombstone collection with policy-based logic          pkg/vm/engine/disttae/txn_table.go Added a new parameter policy to the CollectTombstones method. Implemented logic to handle different tombstone collection policies. Enhanced tombstone collection to include both committed and uncommitted tombstones. +67/-55  txn_table_sharding.go Implement sharding support for tombstone collection            pkg/vm/engine/disttae/txn_table_sharding.go Implemented CollectTombstones with sharding support. Added logic to merge local and remote tombstones. Introduced shardingReader struct and methods. +102/-3  txn_table_sharding_handle.go Add handlers for sharding read operations                                pkg/vm/engine/disttae/txn_table_sharding_handle.go Added handlers for sharding read operations including ReadNext, ReadClose, and ReadCollectTombstones. Implemented logic to handle tombstone collection in sharded environments. +62/-6    server.go Update shard service with new read operations                        pkg/cnservice/server.go Updated shard service initialization to include new read operations. Added handlers for ReadNext, ReadClose, and ReadCollectTombstones. +6/-2      compile.go Use new tombstone collection policy in compile                      pkg/sql/compile/compile.go Updated collectTombstones function to use new tombstone collection policy. +3/-3      types.go Add tombstone collection policy type and constants              pkg/vm/engine/types.go Introduced TombstoneCollectPolicy type and constants. Updated CollectTombstones method signature to include policy parameter. +12/-1    shard.proto Extend proto definitions for sharding operations                  proto/shard.proto Added new message types for tombstone collection and reader operations. Updated ReadParam to include new parameters for sharding operations. +36/-7
Additional files (token-limit)	1 files shard.pb.go ...                                                                                                            pkg/pb/shard/shard.pb.go ... +1270/-262

---

💡 PR-Agent usage:
Comment /help on the PR to get a list of all available PR-Agent tools and their descriptions

[codiumai-pr-agent-pro]

PR-Agent was enabled for this repository. To continue using it, please link your git user with your CodiumAI identity here.

PR Reviewer Guide 🔍

⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪

🧪 No relevant tests

🔒 No security concerns identified

⚡ Key issues to review Code Complexity The CollectTombstones function has become more complex with the addition of policy-based collection. Consider refactoring this function to improve readability and maintainability. Incomplete Implementation The shardingReader struct and its methods (Read, Close, SetOrderBy, GetOrderBy, SetFilterZM) are not fully implemented. These need to be completed to ensure proper functionality. Unimplemented Functions HandleShardingReadNext and HandleShardingReadClose functions are not implemented. These need to be completed to ensure proper functionality of sharding operations.

[codiumai-pr-agent-pro]

PR-Agent was enabled for this repository. To continue using it, please link your git user with your CodiumAI identity here.

PR Code Suggestions ✨

Category	Suggestion	Score
Enhancement	✅ Implement the Read method to handle both local and remote data fetching in the sharding reader Suggestion Impact: The suggestion impacted the commit by implementing the Read method to handle reading from local and remote shards. The commit includes logic to read from a local reader and make remote calls, aligning with the suggestion's intent. code diff: func (r *shardingReader) Read( @@ -398,9 +409,69 @@ vp engine.VectorPool, bat *batch.Batch, ) (isEnd bool, err error) { - - //TODO:: - return false, nil + defer func() { + if err != nil || isEnd { + r.lrd.Close() + //TODO:: close remote reader asynchronously through cleaner. + r.iteratePhase = InEnd + } + }() + for { + switch r.iteratePhase { + case InLocal: + isEnd, err = r.lrd.Read(ctx, cols, expr, mp, vp, bat) + if err != nil { + return + } + if !isEnd { + return + } + relData, err := r.remoteRelData.MarshalBinary() + if err != nil { + return false, err + } + err = r.tblDelegate.forwardRead( + ctx, + shardservice.ReadBuildReader, + func(param *shard.ReadParam) { + param.ReaderBuildParam.RelData = relData + param.ReaderBuildParam.Expr = expr + param.ReaderBuildParam.ReadPolicy = int32(r.remoteReadPolicy) + param.ReaderBuildParam.ScanType = int32(r.remoteScanType) + }, + func(resp []byte) { + r.streamID = types.DecodeUuid(resp) + }, + ) + if err != nil { + return false, err + } + r.iteratePhase = InRemote + case InRemote: + var end bool + err = r.tblDelegate.forwardRead( + ctx, + shardservice.ReadNext, + func(param *shard.ReadParam) { + param.ReadNextParam.Uuid = types.EncodeUuid(&r.streamID) + param.ReadNextParam.Columns = cols + }, + func(resp []byte) { + //TODO:: unmarshall resp into batch and isEnd + //bat.UnmarshalBinary(resp) + //end = ? + }, + ) + if err != nil { + return false, err + } + if !end { + return + } + return true, nil + } + } + } Implement the Read method in the shardingReader struct to handle reading data from both local and remote shards. This should include logic to fetch data from the local reader and make RPC calls to fetch data from remote shards. pkg/vm/engine/disttae/txn_table_sharding.go [393-403] func (r *shardingReader) Read( ctx context.Context, cols []string, expr *plan.Expr, mp *mpool.MPool, vp engine.VectorPool, bat *batch.Batch, ) (isEnd bool, err error) { + // Read from local reader + localIsEnd, err := r.lrd.Read(ctx, cols, expr, mp, vp, bat) + if err != nil { + return false, err + } - //TODO:: - return false, nil + // If local read is not end, return + if !localIsEnd { + return false, nil + } + + // Read from remote shards + err = r.tblDelegate.forwardRead( + ctx, + shardservice.ReadNext, + func(param *shard.ReadParam) { + param.ReadNextParam.Uuid = r.streamID[:] + }, + func(resp []byte) { + // Unmarshal and append to bat + // Set isEnd based on remote read result + }, + ) + + return isEnd, err } Apply this suggestion Suggestion importance[1-10]: 9 Why: The suggestion correctly addresses the incomplete implementation of the Read method by providing a comprehensive approach to handle data fetching from both local and remote shards, which is crucial for the functionality of the shardingReader.	9
	Validate the tombstone collect policy before using it in the function In the HandleShardingReadCollectTombstones function, consider adding a check for the validity of param.CollectTombstonesParam.CollectPolicy before using it. This can prevent potential panics or unexpected behavior if an invalid policy is provided. pkg/vm/engine/disttae/txn_table_sharding_handle.go [224-228] -tombstones, err := tbl.CollectTombstones( - ctx, - 0, - engine.TombstoneCollectPolicy(param.CollectTombstonesParam.CollectPolicy), -) +policy := engine.TombstoneCollectPolicy(param.CollectTombstonesParam.CollectPolicy) +if policy != engine.Policy_CollectUncommittedTombstones && + policy != engine.Policy_CollectCommittedTombstones && + policy != engine.Policy_CollectAllTombstones { + return nil, fmt.Errorf("invalid tombstone collect policy: %v", policy) +} +tombstones, err := tbl.CollectTombstones(ctx, 0, policy) Apply this suggestion Suggestion importance[1-10]: 9 Why: This suggestion is important as it adds validation for the tombstone collect policy, preventing potential runtime errors or panics due to invalid input. Ensuring the validity of input parameters is a critical aspect of robust software development.	9
	✅ Implement the buildShardingReaders function to create and return sharding readers Suggestion Impact: The commit implemented the buildShardingReaders function to create shardingReader instances, including logic to determine local and remote blocks, similar to the suggestion. code diff: func buildShardingReaders( ctx context.Context, - proc any, + p any, expr *plan.Expr, relData engine.RelData, num int, @@ -434,8 +505,103 @@ policy engine.TombstoneApplyPolicy, tbl *txnTableDelegate, ) ([]engine.Reader, error) { - - return nil, nil + var rds []engine.Reader + proc := p.(*process.Process) + + if plan2.IsFalseExpr(expr) { + return []engine.Reader{new(emptyReader)}, nil + } + + if orderBy && num != 1 { + return nil, moerr.NewInternalErrorNoCtx("orderBy only support one reader") + } + + _, uncommittedObjNames := tbl.origin.collectUnCommittedObjects(txnOffset) + uncommittedTombstones, err := tbl.origin.CollectTombstones(ctx, txnOffset, engine.Policy_CollectUncommittedTombstones) + if err != nil { + return nil, err + } + group := func(rd engine.RelData) (local engine.RelData, remote engine.RelData) { + local = rd.BuildEmptyRelData() + remote = rd.BuildEmptyRelData() + engine.ForRangeBlockInfo(0, rd.DataCnt(), rd, func(bi objectio.BlockInfo) (bool, error) { + if bi.IsMemBlk() { + local.AppendBlockInfo(bi) + remote.AppendBlockInfo(bi) + } + if _, ok := uncommittedObjNames[*objectio.ShortName(&bi.BlockID)]; ok { + local.AppendBlockInfo(bi) + } else { + remote.AppendBlockInfo(bi) + } + return true, nil + }) + remote.AttachTombstones(uncommittedTombstones) + return + } + + //relData maybe is nil, indicate that only read data from memory. + if relData == nil || relData.DataCnt() == 0 { + relData = NewEmptyBlockListRelationData() + relData.AppendBlockInfo(objectio.EmptyBlockInfo) + } + + blkCnt := relData.DataCnt() + newNum := num + if blkCnt < num { + newNum = blkCnt + for i := 0; i < num-blkCnt; i++ { + rds = append(rds, new(emptyReader)) + } + } + + scanType := determineScanType(relData, newNum) + mod := blkCnt % newNum + divide := blkCnt / newNum + var shard engine.RelData + var localReadPolicy engine.DataSourceReadPolicy + var remoteReadPolicy engine.DataSourceReadPolicy + for i := 0; i < newNum; i++ { + if i == 0 { + shard = relData.DataSlice(i*divide, (i+1)*divide+mod) + localReadPolicy = engine.Policy_SkipReadCommittedInMem + remoteReadPolicy = engine.Policy_SkipReadUncommittedInMem + } else { + shard = relData.DataSlice(i*divide+mod, (i+1)*divide+mod) + localReadPolicy = engine.Policy_SkipReadInMem + remoteReadPolicy = engine.Policy_SkipReadInMem + } + + localRelData, remoteRelData := group(shard) + ds, err := tbl.origin.buildLocalDataSource(ctx, txnOffset, localRelData, localReadPolicy, policy) + if err != nil { + return nil, err + } + lrd, err := NewReader( + ctx, + proc, + tbl.origin.getTxn().engine, + tbl.origin.GetTableDef(ctx), + tbl.origin.db.op.SnapshotTS(), + expr, + ds, + ) + if err != nil { + return nil, err + } + lrd.scanType = scanType + + srd := &shardingReader{ + lrd: lrd, + tblDelegate: tbl, + remoteRelData: remoteRelData, + remoteReadPolicy: remoteReadPolicy, + remoteScanType: scanType, + } + rds = append(rds, srd) + } + + return rds, nil Implement the buildShardingReaders function to create and return a slice of shardingReader instances. This should include logic to determine which blocks are local and which are remote, and set up the necessary data structures for each reader. pkg/vm/engine/disttae/txn_table_sharding.go [426-439] func buildShardingReaders( ctx context.Context, proc any, expr *plan.Expr, relData engine.RelData, num int, txnOffset int, orderBy bool, policy engine.TombstoneApplyPolicy, tbl *txnTableDelegate, ) ([]engine.Reader, error) { + readers := make([]engine.Reader, 0, num) - return nil, nil + // Build local readers + localReaders, err := tbl.origin.BuildReaders(ctx, proc, expr, relData, num, txnOffset, orderBy) + if err != nil { + return nil, err + } + + // Create sharding readers + for _, lr := range localReaders { + sr := &shardingReader{ + lrd: lr, + tblDelegate: tbl, + streamID: types.NewUuid(), + // TODO: Determine which blocks are remote + } + readers = append(readers, sr) + } + + return readers, nil } Apply this suggestion Suggestion importance[1-10]: 8 Why: The suggestion provides a clear and necessary implementation for the buildShardingReaders function, which is essential for setting up the sharding readers correctly, although it leaves some TODOs for determining remote blocks.	8
	✅ Enhance the Close method to properly terminate both local and remote read operations Suggestion Impact: The commit refactored the Close method to ensure both local and remote read operations are properly terminated, aligning with the suggestion's intention. code diff: func (r *shardingReader) Close() error { - r.lrd.Close() + return r.close() +} + +func (r *shardingReader) close() error { + if !r.closed { + r.lrd.Close() + err := r.tblDelegate.forwardRead( + context.Background(), + shardservice.ReadNext, + func(param *shard.ReadParam) { + param.ReadCloseParam.Uuid = types.EncodeUuid(&r.streamID) + }, + func(resp []byte) { + }, + ) + if err != nil { + return err + } + } return nil Implement the Close method in the shardingReader struct to properly clean up resources, including closing the local reader and terminating any ongoing remote read operations. pkg/vm/engine/disttae/txn_table_sharding.go [406-409] func (r *shardingReader) Close() error { - r.lrd.Close() - return nil + err := r.lrd.Close() + if err != nil { + return err + } + + // Terminate remote read operation + return r.tblDelegate.forwardRead( + context.Background(), + shardservice.ReadClose, + func(param *shard.ReadParam) { + param.ReadCloseParam.Uuid = r.streamID[:] + }, + nil, + ) } Apply this suggestion Suggestion importance[1-10]: 8 Why: The suggestion enhances resource management by ensuring that both local and remote read operations are properly terminated, which is important for preventing resource leaks and ensuring clean shutdowns.	8
	Implement SetOrderBy and GetOrderBy methods to handle result ordering across shards Implement the SetOrderBy and GetOrderBy methods in the shardingReader struct to handle ordering of results across shards. This may involve coordinating with remote shards to ensure correct global ordering. pkg/vm/engine/disttae/txn_table_sharding.go [411-419] func (r *shardingReader) SetOrderBy(orderby []*plan.OrderBySpec) { - //r.source.SetOrderBy(orderby) - panic("not implemented") + r.lrd.SetOrderBy(orderby) + // TODO: Propagate order by to remote shards } func (r *shardingReader) GetOrderBy() []*plan.OrderBySpec { - //return r.source.GetOrderBy() - panic("not implemented") + return r.lrd.GetOrderBy() } Apply this suggestion Suggestion importance[1-10]: 7 Why: The suggestion improves the shardingReader by implementing order handling methods, which enhances the functionality, but it does not fully address the complexity of coordinating order across shards.	7
Possible issue	✅ Implement proper error handling and logging for unimplemented functions Suggestion Impact: The commit added error handling to both HandleShardingReadNext and HandleShardingReadClose functions. It checks for errors and returns them, which aligns with the suggestion to implement proper error handling. code diff: func HandleShardingReadNext( ctx context.Context, shard shard.TableShard, @@ -190,7 +315,76 @@ ts timestamp.Timestamp, buffer *morpc.Buffer, ) ([]byte, error) { - return nil, nil + + tbl, err := getTxnTable( + ctx, + param, + engine, + ) + if err != nil { + return nil, err + } + mp := tbl.proc.Load().Mp() + + streamID := types.DecodeUuid(param.ReadNextParam.Uuid) + cols := param.ReadNextParam.Columns + //find reader by streamID + streamHandler.Lock() + sr, ok := streamHandler.streamReaders[streamID] + if !ok { + streamHandler.Unlock() + return nil, moerr.NewInternalErrorNoCtx("stream reader not found, may be expired") + } + streamHandler.Unlock() + + sr.updateCols(cols, tbl.tableDef) + + buildBatch := func() *batch.Batch { + bat := batch.NewWithSize(len(sr.colTypes)) + bat.Attrs = append(bat.Attrs, cols...) + + for i := 0; i < len(sr.colTypes); i++ { + bat.Vecs[i] = vector.NewVec(sr.colTypes[i]) + + } + return bat + } + bat := buildBatch() + defer func() { + bat.Clean(mp) + }() + + isEnd, err := sr.rd.Read( + ctx, + cols, + nil, + mp, + nil, + bat, + ) + if err != nil { + return nil, err + } + if isEnd { + return buffer.EncodeBytes(types.EncodeBool(&isEnd)), nil + } + + var w bytes.Buffer + if _, err := w.Write(types.EncodeBool(&isEnd)); err != nil { + return nil, err + } + encBat, err := bat.MarshalBinary() + if err != nil { + return nil, err + } + l := uint32(len(encBat)) + if _, err := w.Write(types.EncodeUint32(&l)); err != nil { + return nil, err + } + if _, err := w.Write(encBat); err != nil { + return nil, err + } + return buffer.EncodeBytes(w.Bytes()), nil } func HandleShardingReadClose( @@ -201,6 +395,16 @@ ts timestamp.Timestamp, buffer *morpc.Buffer, ) ([]byte, error) { + streamID := types.DecodeUuid(param.ReadCloseParam.Uuid) + //find reader by streamID + streamHandler.Lock() + defer streamHandler.Unlock() + sr, ok := streamHandler.streamReaders[streamID] + if !ok { + return nil, moerr.NewInternalErrorNoCtx("stream reader not found, may be expired") + } + sr.rd.Close() + delete(streamHandler.streamReaders, sr.streamID) return nil, nil Consider implementing error handling and logging for the HandleShardingReadNext and HandleShardingReadClose functions. Currently, they return nil values without any processing, which might lead to unexpected behavior or silent failures. pkg/vm/engine/disttae/txn_table_sharding_handle.go [185-205] func HandleShardingReadNext( ctx context.Context, shard shard.TableShard, engine engine.Engine, param shard.ReadParam, ts timestamp.Timestamp, buffer *morpc.Buffer, ) ([]byte, error) { - return nil, nil + // TODO: Implement the logic for reading next shard + return nil, fmt.Errorf("HandleShardingReadNext not implemented") } func HandleShardingReadClose( ctx context.Context, shard shard.TableShard, engine engine.Engine, param shard.ReadParam, ts timestamp.Timestamp, buffer *morpc.Buffer, ) ([]byte, error) { - return nil, nil + // TODO: Implement the logic for closing shard read + return nil, fmt.Errorf("HandleShardingReadClose not implemented") } Apply this suggestion Suggestion importance[1-10]: 8 Why: The suggestion correctly identifies that the functions HandleShardingReadNext and HandleShardingReadClose are currently unimplemented and returning nil, which could lead to silent failures. Adding error handling and logging is crucial for identifying and debugging issues, making this a valuable improvement.	8
Maintainability	Extract repeated tombstone collection logic into a separate function to reduce code duplication Consider extracting the repeated code for collecting tombstones into a separate function to improve code reusability and maintainability. This will reduce duplication and make the code easier to update in the future. pkg/sql/compile/compile.go [3778-3813] -tombstone, err = rel.CollectTombstones(ctx, c.TxnOffset, engine.Policy_CollectAllTombstones) +func collectTombstonesWithPolicy(ctx context.Context, relation engine.Relation, txnOffset int) (engine.Tombstoner, error) { + return relation.CollectTombstones(ctx, txnOffset, engine.Policy_CollectAllTombstones) +} + +tombstone, err = collectTombstonesWithPolicy(ctx, rel, c.TxnOffset) if err != nil { return nil, err } ... -subTombstone, err := subrelation.CollectTombstones(ctx, c.TxnOffset, engine.Policy_CollectAllTombstones) +subTombstone, err := collectTombstonesWithPolicy(ctx, subrelation, c.TxnOffset) if err != nil { return nil, err } ... -subTombstone, err := subrelation.CollectTombstones(ctx, c.TxnOffset, engine.Policy_CollectAllTombstones) +subTombstone, err := collectTombstonesWithPolicy(ctx, subrelation, c.TxnOffset) if err != nil { return nil, err } Suggestion importance[1-10]: 7 Why: Extracting repeated code into a separate function enhances code reusability and maintainability. This suggestion effectively reduces duplication, making future updates easier, though it addresses a non-critical aspect of the codebase.	7
Best practice	Use explicit values with iota for enum-like constants to improve clarity and maintainability Consider using the iota operator with explicit values for the TombstoneCollectPolicy constants to make the intent clearer and prevent potential issues if the order of constants changes in the future. pkg/vm/engine/types.go [588-592] const ( - Policy_CollectUncommittedTombstones = 1 << iota - Policy_CollectCommittedTombstones - Policy_CollectAllTombstones = Policy_CollectUncommittedTombstones | Policy_CollectCommittedTombstones + Policy_CollectUncommittedTombstones TombstoneCollectPolicy = 1 << iota + Policy_CollectCommittedTombstones TombstoneCollectPolicy = 1 << iota + Policy_CollectAllTombstones TombstoneCollectPolicy = Policy_CollectUncommittedTombstones | Policy_CollectCommittedTombstones ) Apply this suggestion Suggestion importance[1-10]: 6 Why: While the suggestion to use explicit values with iota for constants is a good practice for clarity and maintainability, the current implementation is already clear. This suggestion offers a minor improvement in terms of future-proofing the code.	6