diff --git a/executor.go b/executor.go index 101e4ec69..5d8c94c75 100644 --- a/executor.go +++ b/executor.go @@ -1511,6 +1511,27 @@ func (e *executor) executeGroupBy(ctx context.Context, index string, c *pql.Call } results, _ := other.([]GroupCount) + // Apply having. + if having, hasHaving, err := c.CallArg("having"); err != nil { + return nil, err + } else if hasHaving { + // parse the condition as PQL + if having.Name != "Condition" { + return nil, errors.New("the only supported having call is Condition()") + } + if len(having.Args) != 1 { + return nil, errors.New("Condition() must contain a single condition") + } + for subj, cond := range having.Args { + switch subj { + case "count", "sum": + results = applyConditionToGroupCounts(results, subj, cond.(*pql.Condition)) + default: + return nil, errors.New("Condition() only supports count or sum") + } + } + } + // Apply offset. if offset, hasOffset, err := c.UintArg("offset"); err != nil { return nil, err @@ -1617,6 +1638,137 @@ func (g GroupCount) Compare(o GroupCount) int { return 0 } +func (g GroupCount) satisfiesCondition(subj string, cond *pql.Condition) bool { + switch subj { + case "count": + switch cond.Op { + case pql.EQ, pql.NEQ, pql.LT, pql.LTE, pql.GT, pql.GTE: + val, ok := cond.Uint64Value() + if !ok { + return false + } + if cond.Op == pql.EQ { + if g.Count == val { + return true + } + } else if cond.Op == pql.NEQ { + if g.Count != val { + return true + } + } else if cond.Op == pql.LT { + if g.Count < val { + return true + } + } else if cond.Op == pql.LTE { + if g.Count <= val { + return true + } + } else if cond.Op == pql.GT { + if g.Count > val { + return true + } + } else if cond.Op == pql.GTE { + if g.Count >= val { + return true + } + } + case pql.BETWEEN, pql.BTWN_LT_LTE, pql.BTWN_LTE_LT, pql.BTWN_LT_LT: + val, ok := cond.Uint64SliceValue() + if !ok { + return false + } + if cond.Op == pql.BETWEEN { + if val[0] <= g.Count && g.Count <= val[1] { + return true + } + } else if cond.Op == pql.BTWN_LT_LTE { + if val[0] < g.Count && g.Count <= val[1] { + return true + } + } else if cond.Op == pql.BTWN_LTE_LT { + if val[0] <= g.Count && g.Count < val[1] { + return true + } + } else if cond.Op == pql.BTWN_LT_LT { + if val[0] < g.Count && g.Count < val[1] { + return true + } + } + } + case "sum": + switch cond.Op { + case pql.EQ, pql.NEQ, pql.LT, pql.LTE, pql.GT, pql.GTE: + val, ok := cond.Int64Value() + if !ok { + return false + } + if cond.Op == pql.EQ { + if g.Sum == val { + return true + } + } else if cond.Op == pql.NEQ { + if g.Sum != val { + return true + } + } else if cond.Op == pql.LT { + if g.Sum < val { + return true + } + } else if cond.Op == pql.LTE { + if g.Sum <= val { + return true + } + } else if cond.Op == pql.GT { + if g.Sum > val { + return true + } + } else if cond.Op == pql.GTE { + if g.Sum >= val { + return true + } + } + case pql.BETWEEN, pql.BTWN_LT_LTE, pql.BTWN_LTE_LT, pql.BTWN_LT_LT: + val, ok := cond.Int64SliceValue() + if !ok { + return false + } + if cond.Op == pql.BETWEEN { + if val[0] <= g.Sum && g.Sum <= val[1] { + return true + } + } else if cond.Op == pql.BTWN_LT_LTE { + if val[0] < g.Sum && g.Sum <= val[1] { + return true + } + } else if cond.Op == pql.BTWN_LTE_LT { + if val[0] <= g.Sum && g.Sum < val[1] { + return true + } + } else if cond.Op == pql.BTWN_LT_LT { + if val[0] < g.Sum && g.Sum < val[1] { + return true + } + } + } + } + return false +} + +// applyConditionToGroupCounts filters the contents of gcs according +// to the condition. Currently, `count` and `sum` are the only +// fields supported. +func applyConditionToGroupCounts(gcs []GroupCount, subj string, cond *pql.Condition) []GroupCount { + var i int + for _, gc := range gcs { + if !gc.satisfiesCondition(subj, cond) { + continue // drop this GroupCount + } + gcs[i] = gc + i++ + } + return gcs[:i] +} + func (e *executor) executeGroupByShard(ctx context.Context, index string, c *pql.Call, filter *pql.Call, shard uint64, childRows []RowIDs) (_ []GroupCount, err error) { span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeGroupByShard") defer span.Finish() diff --git a/executor_internal_test.go b/executor_internal_test.go index e85c9edb1..ee2880509 100644 --- a/executor_internal_test.go +++ b/executor_internal_test.go @@ -46,7 +46,7 @@ func TestExecutor_TranslateGroupByCall(t *testing.T) { t.Fatalf("creating fields %v, %v, %v", erra, errb, errc) } - query, err := pql.ParseString(`GroupBy(Rows(ak), Rows(b), Rows(ck), previous=["la", 0, "ha"])`) + query, err := pql.ParseString(`GroupBy(Rows(ak), Rows(b), Rows(ck), previous=["la", 0, "ha"], having=Condition(count > 10))`) if err != nil { t.Fatalf("parsing query: %v", err) } @@ -64,6 +64,19 @@ func TestExecutor_TranslateGroupByCall(t *testing.T) { } } + if having, hok := c.Args["having"].(*pql.Call); !hok { + t.Fatal("expected having to be a call") + } else if cond, cok := having.Args["count"].(*pql.Condition); !cok { + t.Fatal("expected condition to be a count") + } else if cond.Op != pql.GT { + t.Fatal("expected condition op to be >") + } else { + val, ok := cond.Uint64Value() + if !ok || val != uint64(10) { + t.Fatal("expected condition val to be uint64(10)") + } + } + errTests := []struct { pql string err string @@ -222,3 +235,123 @@ func TestFieldRowMarshalJSON(t *testing.T) { t.Fatalf("unexpected json: %s", b) } } + +func TestExecutor_GroupCountCondition(t *testing.T) { + t.Run("satisfiesCondition", func(t *testing.T) { + type condCheck struct { + cond string + exp bool + } + tests := []struct { + groupCount GroupCount + checks []condCheck + }{ + { + groupCount: GroupCount{Count: 100}, + checks: []condCheck{ + {cond: "count == 99", exp: false}, + {cond: "count != 99", exp: true}, + {cond: "count < 99", exp: false}, + {cond: "count <= 99", exp: false}, + {cond: "count > 99", exp: true}, + {cond: "count >= 99", exp: true}, + + {cond: "count == 100", exp: true}, + {cond: "count != 100", exp: false}, + {cond: "count < 100", exp: false}, + {cond: "count <= 100", exp: true}, + {cond: "count > 100", exp: false}, + {cond: "count >= 100", exp: true}, + + {cond: "count == 101", exp: false}, + {cond: "count != 101", exp: true}, + {cond: "count < 101", exp: true}, + {cond: "count <= 101", exp: true}, + {cond: "count > 101", exp: false}, + {cond: "count >= 101", exp: false}, + }, + }, + { + groupCount: GroupCount{Sum: 100}, + checks: []condCheck{ + {cond: "sum == 99", exp: false}, + {cond: "sum != 99", exp: true}, + {cond: "sum < 99", exp: false}, + {cond: "sum <= 99", exp: false}, + {cond: "sum > 99", exp: true}, + {cond: "sum >= 99", exp: true}, + + {cond: "sum == 100", exp: true}, + {cond: "sum != 100", exp: false}, + {cond: "sum < 100", exp: false}, + {cond: "sum <= 100", exp: true}, + {cond: "sum > 100", exp: false}, + {cond: "sum >= 100", exp: true}, + + {cond: "sum == 101", exp: false}, + {cond: "sum != 101", exp: true}, + {cond: "sum < 101", exp: true}, + {cond: "sum <= 101", exp: true}, + {cond: "sum > 101", exp: false}, + {cond: "sum >= 101", exp: false}, + }, + }, + { + groupCount: GroupCount{Sum: -100}, + checks: []condCheck{ + {cond: "sum == -99", exp: false}, + {cond: "sum != -99", exp: true}, + {cond: "sum < -99", exp: true}, + {cond: "sum <= -99", exp: true}, + {cond: "sum > -99", exp: false}, + {cond: "sum >= -99", exp: false}, + + {cond: "sum == -100", exp: true}, + {cond: "sum != -100", exp: false}, + {cond: "sum < -100", exp: false}, + {cond: "sum <= -100", exp: true}, + {cond: "sum > -100", exp: false}, + {cond: "sum >= -100", exp: true}, + + {cond: "sum == -101", exp: false}, + {cond: "sum != -101", exp: true}, + {cond: "sum < -101", exp: false}, + {cond: "sum <= -101", exp: false}, + {cond: "sum > -101", exp: true}, + {cond: "sum >= -101", exp: true}, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test (#%d):", i), func(t *testing.T) { + for j, check := range test.checks { + t.Run(fmt.Sprintf("check (#%d):", j), func(t *testing.T) { + + query, err := pql.ParseString(fmt.Sprintf("GroupBy(Rows(a), having=Condition(%s))", check.cond)) + if err != nil { + t.Fatalf("parsing query: %v", err) + } + c := query.Calls[0] + having := c.Args["having"].(*pql.Call) + + var got bool + for subj, cond := range having.Args { + switch subj { + case "count", "sum": + condition, ok := cond.(*pql.Condition) + if !ok { + t.Fatalf("not a valid condition") + } + got = test.groupCount.satisfiesCondition(subj, condition) + } + } + + if got != check.exp { + t.Fatalf("expected: %v, but got: %v", check.exp, got) + } + }) + } + }) + } + }) +} diff --git a/executor_test.go b/executor_test.go index 294a3f256..212b1eebf 100644 --- a/executor_test.go +++ b/executor_test.go @@ -3459,6 +3459,22 @@ func TestExecutor_GroupByStrings(t *testing.T) { {Group: []pilosa.FieldRow{{Field: "generals", RowID: 2, RowKey: "r2"}}, Count: 5, Sum: 30}, }, }, + { + query: "GroupBy(Rows(generals), aggregate=Sum(field=v), having=Condition(sum>25))", + expected: []pilosa.GroupCount{ + {Group: []pilosa.FieldRow{{Field: "generals", RowID: 2, RowKey: "r2"}}, Count: 5, Sum: 30}, + }, + }, + { + query: "GroupBy(Rows(generals), aggregate=Sum(field=v), having=Condition(-55))", + expected: []pilosa.GroupCount{}, + }, } for i, tst := range tests { diff --git a/pql/ast.go b/pql/ast.go index 0476a868a..4d828ae22 100644 --- a/pql/ast.go +++ b/pql/ast.go @@ -387,6 +387,7 @@ var callInfoByFunc = map[string]callInfo{ "limit": int64(0), "previous": nil, "aggregate": nil, + "having": nil, }, }, "Options": { @@ -777,6 +778,79 @@ func (cond *Condition) IntSliceValue() ([]int64, error) { } } +func (cond *Condition) Uint64Value() (uint64, bool) { + val := cond.Value + + switch tval := val.(type) { + case int64: + if tval >= 0 { + return uint64(tval), true + } + case uint64: + return tval, true + } + + return 0, false +} + +func (cond *Condition) Uint64SliceValue() ([]uint64, bool) { + val := cond.Value + + switch tval := val.(type) { + case []interface{}: + ret := make([]uint64, len(tval)) + for i, v := range tval { + switch tv := v.(type) { + case int64: + ret[i] = uint64(tv) + case uint64: + ret[i] = tv + default: + return nil, false + } + } + return ret, true + } + + return nil, false +} + +func (cond *Condition) Int64Value() (int64, bool) { + val := cond.Value + + switch tval := val.(type) { + case int64: + return tval, true + case uint64: + // TODO: consider overflow? + return int64(tval), true + } + + return 0, false +} + +func (cond *Condition) Int64SliceValue() ([]int64, bool) { + val := cond.Value + + switch tval := val.(type) { + case []interface{}: + ret := make([]int64, len(tval)) + for i, v := range tval { + switch tv := v.(type) { + case int64: + ret[i] = tv + case uint64: + ret[i] = int64(tv) + default: + return nil, false + } + } + return ret, true + } + + return nil, false +} + func formatValue(v interface{}) string { switch v := v.(type) { case string: