feat(window): support statewindow partition by (#3936)

Signed-off-by: Song Gao <disxiaofei@163.com>

Author: Yisaer

.github/workflows/markdown_config.json

@@ -10,5 +10,6 @@
     "MD045": false,
     "MD046": false,
     "MD051": false,
-    "MD059": false
+    "MD059": false,
+    "MD060": false
   }

docs/en_US/sqls/windows.md

@@ -112,6 +112,30 @@ Output:
 [{"a":2},{"a":2}]
 ```
 
+#### State Window Partitioning
+
+We can perform partitioning calculations on a state window using the `partition by` clause, as follows:
+
+```sql
+SELECT * from demo group by statewindow(a = 1, a = 5) over (partition by b)
+```
+
+For the following input:
+
+```txt
+{"a":1,"b":1}
+{"a":1,"b":2}
+{"a":5,"b":1}
+```
+
+The output is as follows:
+
+```txt
+[{"a":1,"b":1},{"a":5,"b":1}]
+```
+
+The partition `b=2` didn't be output due to the partition haven't trigger the condition yet.
+
 ## Count window
 
 Please notice that the count window does not concern time, it only concern about events count.

docs/zh_CN/sqls/windows.md

@@ -114,6 +114,30 @@ SELECT * from demo group by statewindow(a > 1)
 [{"a":2},{"a":2}]
 ```
 
+#### 条件窗口分区
+
+我们可以通过 partition by 子句对条件窗口进行分区计算，如下:
+
+```sql
+SELECT * from demo group by statewindow(a = 1, a = 5) over (partition by b)
+```
+
+此时对于如下输入:
+
+```txt
+{"a":1,"b":1}
+{"a":1,"b":2}
+{"a":5,"b":1}
+```
+
+输出如下:
+
+```txt
+[{"a":1,"b":1},{"a":5,"b":1}]
+```
+
+`b=2` 的数据并没有被输出，因为该分区窗口并未满足条件。
+
 ## 计数窗口
 
 请注意计数窗口不关注时间，只关注事件发生的次数。

internal/topo/node/window_op.go

@@ -56,7 +56,8 @@ type WindowConfig struct {
 	SingleCondition ast.Expr
 	BeginCondition  ast.Expr
 	EmitCondition   ast.Expr
-	Dimensions      ast.Dimensions
+
+	PartitionExpr *ast.PartitionExpr
 }
 
 type WindowOperator struct {

internal/topo/node/window_v2_op.go

@@ -39,6 +39,7 @@ func init() {
 	gob.Register(&WindowScanner{})
 	gob.Register(time.Time{})
 	gob.Register(&StateWindowStatus{})
+	gob.Register(map[string]*StateWindowStatus{})
 }
 
 type WindowV2Operator struct {
@@ -106,7 +107,7 @@ type WindowV2Exec interface {
 type StateWindowOp struct {
 	*WindowV2Operator
 	status          map[string]*StateWindowStatus
-	Dimensions      ast.Dimensions
+	PartitionExpr   *ast.PartitionExpr
 	SingleCondition ast.Expr
 	BeginCondition  ast.Expr
 	EmitCondition   ast.Expr
@@ -127,8 +128,8 @@ func NewStateWindowOp(o *WindowV2Operator) *StateWindowOp {
 		EmitCondition:    o.windowConfig.EmitCondition,
 		SingleCondition:  o.windowConfig.SingleCondition,
 		stateFuncs:       o.windowConfig.StateFuncs,
-		Dimensions:       o.windowConfig.Dimensions,
 		status:           make(map[string]*StateWindowStatus),
+		PartitionExpr:    o.windowConfig.PartitionExpr,
 	}
 }
 
@@ -145,6 +146,23 @@ func (s *StateWindowOp) emit(ctx api.StreamContext, status *StateWindowStatus) {
 	s.onSend(ctx, results)
 }
 
+func calPartition(fv *xsql.FunctionValuer, partitionExpr *ast.PartitionExpr, row *xsql.Tuple) string {
+	name := "parKey_"
+	if partitionExpr == nil {
+		return name
+	}
+	ve := &xsql.ValuerEval{Valuer: xsql.MultiValuer(row, fv, &xsql.WildcardValuer{Data: row})}
+	for _, expr := range partitionExpr.Exprs {
+		r := ve.Eval(expr)
+		if _, ok := r.(error); ok {
+			continue
+		} else {
+			name += fmt.Sprintf("%v,", r)
+		}
+	}
+	return name
+}
+
 func (s *StateWindowOp) exec(ctx api.StreamContext, errCh chan<- error) {
 	v, err := ctx.GetState(V2WindowInputsKey)
 	if err == nil && v != nil {
@@ -166,7 +184,7 @@ func (s *StateWindowOp) exec(ctx api.StreamContext, errCh chan<- error) {
 			s.onProcessStart(ctx, input)
 			switch row := data.(type) {
 			case *xsql.Tuple:
-				name := calDimension(fv, s.Dimensions, row)
+				name := calPartition(fv, s.PartitionExpr, row)
 				status, ok := s.status[name]
 				if !ok {
 					status = &StateWindowStatus{

internal/topo/node/window_v2_op_test.go

@@ -37,6 +37,65 @@ func init() {
 	testx.InitEnv("node_test")
 }
 
+func TestStateWindowStatePartition(t *testing.T) {
+	conf.IsTesting = true
+	ctx, cancel := mockContext.NewMockContext("1", "2").WithCancel()
+	now := time.Now()
+	o := &def.RuleOption{
+		BufferLength: 10,
+	}
+	kv, err := store.GetKV("stream")
+	require.NoError(t, err)
+	require.NoError(t, prepareStream())
+	sql := "select count(*) from stream group by statewindow(a =1 , a = 2) over (partition by b)"
+	stmt, err := xsql.NewParser(strings.NewReader(sql)).Parse()
+	require.NoError(t, err)
+	p, err := planner.CreateLogicalPlan(stmt, o, kv)
+	require.NoError(t, err)
+	require.NotNil(t, p)
+	windowPlan := extractWindowPlan(p)
+	require.NotNil(t, windowPlan)
+	op, err := node.NewWindowV2Op("window", node.WindowConfig{
+		Type:           windowPlan.WindowType(),
+		BeginCondition: windowPlan.GetBeginCondition(),
+		EmitCondition:  windowPlan.GetEmitCondition(),
+		PartitionExpr:  windowPlan.GetPartitionExpr(),
+	}, o)
+	require.NoError(t, err)
+	require.NotNil(t, op)
+	input, _ := op.GetInput()
+	output := make(chan any, 10)
+	op.AddOutput(output, "output")
+	errCh := make(chan error, 10)
+	op.Exec(ctx, errCh)
+	waitExecute()
+	input <- &xsql.Tuple{Message: map[string]any{"a": true}, Timestamp: now}
+	input <- &xsql.Tuple{Message: map[string]any{"a": false}, Timestamp: now.Add(500 * time.Millisecond)}
+	waitExecute()
+	input <- &xsql.Tuple{Message: map[string]any{"a": int64(1), "b": int64(1)}, Timestamp: now}
+	input <- &xsql.Tuple{Message: map[string]any{"a": int64(1), "b": int64(2)}, Timestamp: now}
+	input <- &xsql.Tuple{Message: map[string]any{"a": int64(2), "b": int64(1)}, Timestamp: now.Add(500 * time.Millisecond)}
+	waitExecute()
+	got := <-output
+	wt, ok := got.(*xsql.WindowTuples)
+	require.True(t, ok)
+	require.NotNil(t, wt)
+	d := wt.ToMaps()
+	require.Equal(t, []map[string]any{
+		{
+			"a": int64(1),
+			"b": int64(1),
+		},
+		{
+			"a": int64(2),
+			"b": int64(1),
+		},
+	}, d)
+	cancel()
+	waitExecute()
+	op.Close()
+}
+
 func TestStateWindowState(t *testing.T) {
 	conf.IsTesting = true
 	ctx, cancel := mockContext.NewMockContext("1", "2").WithCancel()

internal/topo/planner/planner.go

@@ -423,7 +423,7 @@ func buildOps(lp LogicalPlan, tp *topo.Topo, options *def.RuleOption, sources ma
 			BeginCondition:   t.beginCondition,
 			EmitCondition:    t.emitCondition,
 			SingleCondition:  t.singleCondition,
-			Dimensions:       t.dimensions,
+			PartitionExpr:    t.PartitionExpr,
 			StateFuncs:       t.stateFuncs,
 		}
 		// state window only support v2 window
@@ -644,7 +644,7 @@ func createLogicalPlanFull(stmt *ast.SelectStatement, opt *def.RuleOption, store
 					beginCondition:  w.BeginCondition,
 					emitCondition:   w.EmitCondition,
 					singleCondition: w.SingleCondition,
-					dimensions:      dimensions.GetGroups(),
+					PartitionExpr:   w.PartitionExpr,
 				}.Init()
 				if w.Length != nil {
 					wp.length = int(w.Length.Val)

internal/topo/planner/windowPlan.go

@@ -24,7 +24,7 @@ import (
 
 type WindowPlan struct {
 	baseLogicalPlan
-	dimensions       ast.Dimensions
+	PartitionExpr    *ast.PartitionExpr
 	singleCondition  ast.Expr
 	beginCondition   ast.Expr
 	emitCondition    ast.Expr
@@ -67,6 +67,10 @@ func (p *WindowPlan) GetEmitCondition() ast.Expr {
 	return p.emitCondition
 }
 
+func (p *WindowPlan) GetPartitionExpr() *ast.PartitionExpr {
+	return p.PartitionExpr
+}
+
 func (p *WindowPlan) BuildExplainInfo() {
 	t := p.wtype.String()
 	info := "{ length:" + strconv.Itoa(p.length) + ", "

internal/xsql/parser.go

@@ -972,13 +972,10 @@ func (p *Parser) parseCall(n string) (ast.Expr, error) {
 			win.Filter = f
 		}
 		// parse over when clause
-		c, err := p.ParseOver4Window()
+		err = p.ParseOver4Window(win)
 		if err != nil {
 			return nil, err
-		} else if c != nil {
-			win.TriggerCondition = c
 		}
-
 		return win, nil
 	}
 }
@@ -1616,25 +1613,28 @@ func (p *Parser) parseStreamOptions() (*ast.Options, error) {
 	return opts, nil
 }
 
-func (p *Parser) ParseOver4Window() (ast.Expr, error) {
+func (p *Parser) ParseOver4Window(win *ast.Window) error {
 	if tok, _ := p.scanIgnoreWhitespace(); tok != ast.OVER {
 		p.unscan()
-		return nil, nil
+		return nil
 	}
 	if tok, lit := p.scanIgnoreWhitespace(); tok != ast.LPAREN {
-		return nil, fmt.Errorf("Found %q after OVER, expect parentheses.", lit)
+		return fmt.Errorf("Found %q after OVER, expect parentheses.", lit)
 	}
-	if tok, lit := p.scanIgnoreWhitespace(); tok != ast.WHEN {
-		return nil, fmt.Errorf("Found %q after OVER(, expect WHEN.", lit)
+	var err error
+	win.TriggerCondition, err = p.parseWhen()
+	if err != nil {
+		return err
 	}
-	expr, err := p.ParseExpr()
+	win.PartitionExpr, err = p.parsePartitionBy()
 	if err != nil {
-		return nil, err
+		return err
 	}
+
 	if tok, lit := p.scanIgnoreWhitespace(); tok != ast.RPAREN {
-		return nil, fmt.Errorf("Found %q after OVER, expect right parentheses.", lit)
+		return fmt.Errorf("Found %q after OVER, expect right parentheses.", lit)
 	}
-	return expr, nil
+	return nil
 }
 
 // Only support filter on window now

pkg/ast/statement.go

@@ -205,6 +205,7 @@ func (w WindowType) String() string {
 }
 
 type Window struct {
+	PartitionExpr    *PartitionExpr
 	TriggerCondition Expr
 	SingleCondition  Expr
 	BeginCondition   Expr