128 lines
2.6 KiB
Go
128 lines
2.6 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"time"
|
|
|
|
"git.kingecg.top/kingecg/gomog/internal/engine"
|
|
"git.kingecg.top/kingecg/gomog/pkg/types"
|
|
)
|
|
|
|
func main() {
|
|
// 创建内存存储和流式聚合引擎
|
|
store := engine.NewMemoryStore(nil)
|
|
aggEngine := engine.NewStreamAggregationEngine(store)
|
|
|
|
// 创建测试数据
|
|
collection := "test_stream"
|
|
testDocs := make(map[string]types.Document)
|
|
for i := 0; i < 1000; i++ {
|
|
testDocs[fmt.Sprintf("doc%d", i)] = types.Document{
|
|
ID: fmt.Sprintf("doc%d", i),
|
|
Data: map[string]interface{}{
|
|
"name": fmt.Sprintf("User%d", i),
|
|
"age": 20 + (i % 50),
|
|
"score": float64(50 + (i % 50)),
|
|
"status": map[string]interface{}{"active": i%2 == 0},
|
|
},
|
|
}
|
|
}
|
|
|
|
// 插入测试数据
|
|
var docs []types.Document
|
|
for _, doc := range testDocs {
|
|
docs = append(docs, doc)
|
|
}
|
|
|
|
if err := store.InsertMany(collection, docs); err != nil {
|
|
log.Printf("Error inserting documents: %v", err)
|
|
return
|
|
}
|
|
|
|
// 定义聚合管道
|
|
pipeline := []types.AggregateStage{
|
|
{
|
|
Stage: "$match",
|
|
Spec: map[string]interface{}{
|
|
"age": map[string]interface{}{
|
|
"$gte": 25,
|
|
"$lte": 35,
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Stage: "$project",
|
|
Spec: map[string]interface{}{
|
|
"name": 1,
|
|
"age": 1,
|
|
"score": 1,
|
|
},
|
|
},
|
|
{
|
|
Stage: "$addFields",
|
|
Spec: map[string]interface{}{
|
|
"isHighScorer": map[string]interface{}{
|
|
"$gte": []interface{}{"$score", 80.0},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Stage: "$sort",
|
|
Spec: map[string]interface{}{
|
|
"score": -1,
|
|
},
|
|
},
|
|
{
|
|
Stage: "$limit",
|
|
Spec: 10,
|
|
},
|
|
}
|
|
|
|
// 执行流式聚合
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
opts := engine.StreamAggregationOptions{
|
|
BufferSize: 100, // 每次处理100个文档
|
|
}
|
|
|
|
resultChan, errChan := aggEngine.StreamExecute(ctx, collection, pipeline, opts)
|
|
|
|
// 处理结果
|
|
var finalResults []types.Document
|
|
for {
|
|
select {
|
|
case batch, ok := <-resultChan:
|
|
if !ok {
|
|
resultChan = nil
|
|
continue
|
|
}
|
|
fmt.Printf("Received batch with %d documents\n", len(batch))
|
|
finalResults = append(finalResults, batch...)
|
|
case err, ok := <-errChan:
|
|
if ok && err != nil {
|
|
log.Printf("Error during stream aggregation: %v", err)
|
|
return
|
|
}
|
|
if !ok {
|
|
// 通道已关闭，结束循环
|
|
goto done
|
|
}
|
|
case <-ctx.Done():
|
|
log.Println("Context cancelled")
|
|
return
|
|
}
|
|
}
|
|
|
|
done:
|
|
fmt.Printf("Total results: %d\n", len(finalResults))
|
|
for i, doc := range finalResults {
|
|
if i < 5 { // 只打印前5个结果
|
|
fmt.Printf("Result %d: ID=%s, Name=%s, Age=%v, Score=%v\n",
|
|
i+1, doc.ID, doc.Data["name"], doc.Data["age"], doc.Data["score"])
|
|
}
|
|
}
|
|
}
|