Repository: polarsignals/frostdb (go)

Issue:
Support more aggregation functions
Right now frostdb only has limited support for aggregation functions. For observability and analytical purpose, more aggregation functions should be supported to cover more use cases. For me, I have the requirement to support aggregations like `count`, `average` and `quantile`. 

I opened this as an umbrella issue to track the support to other aggregation functions. For aggregation functions to support, feel free to add your use case. For a good start, probably we can take a look at https://prometheus.io/docs/prometheus/latest/querying/operators/#aggregation-operators and https://duckdb.org/docs/sql/aggregates.

- [x] Support `count` aggregate function. PR: https://github.com/polarsignals/frostdb/pull/176
- [ ] Support `quantile` aggregate function
- [ ] Support `average` aggregate function



Pull request description:
Add support for average aggregations
Thanks to @asubiotto we had the idea to make the average aggregation a projection. 

The way it works now is that the logical plan creates an average aggregation but then optimizes those into the two `sum` and `count` aggregations plus an `avg` projection. The logical plan makes sure to wrap both aggregations with the average projection, such that anything reading the column afterward will only see the final result. 

TODO:
- [x] Improve `AverageAggregationPushDown.Optimize()` to support with and without alias columns

Closes #177 
Superseeds #266 & #190

Patch:
diff --git a/query/logicalplan/expr.go b/query/logicalplan/expr.go
index 007537a..2ef2af9 100644
--- a/query/logicalplan/expr.go
+++ b/query/logicalplan/expr.go
@@ -450,6 +450,7 @@ const (
 	AggFuncMin
 	AggFuncMax
 	AggFuncCount
+	AggFuncAvg
 )
 
 func (f AggFunc) String() string {
@@ -462,6 +463,8 @@ func (f AggFunc) String() string {
 		return "max"
 	case AggFuncCount:
 		return "count"
+	case AggFuncAvg:
+		return "avg"
 	default:
 		panic("unknown aggregation function")
 	}
@@ -495,6 +498,13 @@ func Count(expr Expr) *AggregationFunction {
 	}
 }
 
+func Avg(expr Expr) *AggregationFunction {
+	return &AggregationFunction{
+		Func: AggFuncAvg,
+		Expr: expr,
+	}
+}
+
 type AliasExpr struct {
 	Expr  Expr
 	Alias string
@@ -590,3 +600,40 @@ func (d *DurationExpr) Computed() bool {
 func (d *DurationExpr) Value() time.Duration {
 	return d.duration
 }
+
+type AverageExpr struct {
+	Expr Expr
+}
+
+func (a *AverageExpr) DataType(s *parquet.Schema) (arrow.DataType, error) {
+	return a.Expr.DataType(s)
+}
+
+func (a *AverageExpr) Name() string {
+	return a.Expr.Name()
+}
+
+func (a *AverageExpr) ColumnsUsedExprs() []Expr {
+	return a.Expr.ColumnsUsedExprs()
+}
+
+func (a *AverageExpr) MatchPath(path string) bool {
+	return a.Expr.MatchPath(path)
+}
+
+func (a *AverageExpr) MatchColumn(name string) bool {
+	return a.Expr.MatchColumn(name)
+}
+
+func (a *AverageExpr) Computed() bool {
+	return true
+}
+
+func (a *AverageExpr) Accept(visitor Visitor) bool {
+	continu := visitor.PreVisit(a)
+	if !continu {
+		return false
+	}
+
+	return visitor.PostVisit(a)
+}
diff --git a/query/logicalplan/optimize.go b/query/logicalplan/optimize.go
index 68a0ded..438706c 100644
--- a/query/logicalplan/optimize.go
+++ b/query/logicalplan/optimize.go
@@ -1,16 +1,81 @@
 package logicalplan
 
+import (
+	"golang.org/x/exp/slices"
+)
+
 type Optimizer interface {
 	Optimize(plan *LogicalPlan) *LogicalPlan
 }
 
 var DefaultOptimizers = []Optimizer{
+	&AverageAggregationPushDown{},
 	&PhysicalProjectionPushDown{},
 	&FilterPushDown{},
 	&DistinctPushDown{},
 	&ProjectionPushDown{},
 }
 
+type AverageAggregationPushDown struct{}
+
+func (p *AverageAggregationPushDown) Optimize(plan *LogicalPlan) *LogicalPlan {
+	if plan.Aggregation == nil {
+		return plan
+	}
+
+	for i, aggExpr := range plan.Aggregation.AggExprs {
+		var aggFunc *AggregationFunction
+		var alias *AliasExpr
+		var column Expr
+
+		// In case the aggregation contains an alias
+		if aliasExpr, ok := aggExpr.(*AliasExpr); ok {
+			alias = aliasExpr
+			if af, ok := aliasExpr.Expr.(*AggregationFunction); ok {
+				if af.Func == AggFuncAvg {
+					column = af.Expr
+					aggFunc = af
+				}
+			}
+		}
+		if ae, ok := aggExpr.(*AggregationFunction); ok {
+			if ae.Func == AggFuncAvg {
+				column = ae.Expr
+				aggFunc = ae
+			}
+		}
+
+		if aggFunc == nil {
+			// no aggregation func found, skipping
+			continue
+		}
+
+		// Delete this average aggregation from the logicalplan.
+		plan.Aggregation.AggExprs = slices.Delete(plan.Aggregation.AggExprs, i, i+1)
+		// Add sum and count aggregation for the column to the logicalplan.
+		plan.Aggregation.AggExprs = append(plan.Aggregation.AggExprs,
+			Sum(aggFunc.Expr),
+			Count(aggFunc.Expr),
+		)
+
+		projection := &AverageExpr{Expr: column}
+		if alias != nil {
+			alias.Expr = column
+			projection.Expr = alias
+		}
+
+		// Wrap the aggregations with the average projection to always call it after aggregating.
+		plan = &LogicalPlan{
+			Input: plan,
+			Projection: &Projection{
+				Exprs: []Expr{projection},
+			},
+		}
+	}
+
+	return plan
+}
+
 // The PhysicalProjectionPushDown optimizer tries to push down the actual
 // physical columns used by the query to the table scan, so the table provider
 // can decide to only read the columns that are actually going to be used by
diff --git a/query/physicalplan/project.go b/query/physicalplan/project.go
index 0bd3d1a..23a2110 100644
--- a/query/physicalplan/project.go
+++ b/query/physicalplan/project.go
@@ -184,6 +184,8 @@ func projectionFromExpr(expr logicalplan.Expr) (columnProjection, error) {
 		return binaryExprProjection{
 			boolExpr: boolExpr,
 		}, nil
+	case *logicalplan.AverageExpr:
+		return &averageProjection{expr: e}, nil
 	default:
 		return nil, fmt.Errorf("unsupported expression type for projection: %T", expr)
 	}
@@ -271,3 +273,71 @@ func (p *Projection) Draw() *Diagram {
 	details := fmt.Sprintf("Projection (%s)", strings.Join(columns, ","))
 	return &Diagram{Details: details, Child: child}
 }
+
+type averageProjection struct {
+	expr logicalplan.Expr
+}
+
+func (a *averageProjection) Name() string {
+	return a.expr.Name()
+}
+
+func (a *averageProjection) Project(mem memory.Allocator, r arrow.Record) ([]arrow.Field, []arrow.Array, error) {
+	columnName := a.expr.Name()
+	resultName := "avg(" + columnName + ")"
+	if avgExpr, ok := a.expr.(*logicalplan.AverageExpr); ok {
+		if ae, ok := avgExpr.Expr.(*logicalplan.AliasExpr); ok {
+			columnName = ae.Expr.Name()
+			resultName = ae.Alias
+		}
+	}
+
+	columnSum := "sum(" + columnName + ")"
+	columnCount := "count(" + columnName + ")"
+
+	schema := r.Schema()
+
+	sumIndex := schema.FieldIndices(columnSum)
+	if len(sumIndex) != 1 {
+		return nil, nil, fmt.Errorf("sum column for average projection for column %s not found", columnName)
+	}
+	countIndex := schema.FieldIndices(columnCount)
+	if len(countIndex) != 1 {
+		return nil, nil, fmt.Errorf("count column for average projection for column %s not found", columnName)
+	}
+
+	sums := r.Column(sumIndex[0])
+	counts := r.Column(countIndex[0])
+
+	fields := make([]arrow.Field, 0, len(schema.Fields())-1)
+	columns := make([]arrow.Array, 0, len(schema.Fields())-1)
+
+	// Only add the fields and columns that aren't the average's underlying sum and count columns.
+	for i, field := range schema.Fields() {
+		if i != sumIndex[0] && i != countIndex[0] {
+			fields = append(fields, field)
+			columns = append(columns, r.Column(i))
+		}
+	}
+
+	// Add the field and column for the projected average aggregation.
+	fields = append(fields, arrow.Field{
+		Name: resultName,
+		Type: &arrow.Int64Type{},
+	})
+	columns = append(columns, avgInt64arrays(mem, sums, counts))
+
+	return fields, columns, nil
+}
+
+func avgInt64arrays(pool memory.Allocator, sums, counts arrow.Array) arrow.Array {
+	sumsInts := sums.(*array.Int64)
+	countsInts := counts.(*array.Int64)
+
+	res := array.NewInt64Builder(pool)
+	for i := 0; i < sumsInts.Len(); i++ {
+		res.Append(sumsInts.Value(i) / countsInts.Value(i))
+	}
+
+	return res.NewArray()
+}
diff --git a/sqlparse/visitor.go b/sqlparse/visitor.go
index e550ed2..04a4d33 100644
--- a/sqlparse/visitor.go
+++ b/sqlparse/visitor.go
@@ -90,6 +90,8 @@ func (v *astVisitor) leaveImpl(n ast.Node) error {
 			v.exprStack[lastExpr] = logicalplan.Min(v.exprStack[lastExpr])
 		case "max":
 			v.exprStack[lastExpr] = logicalplan.Max(v.exprStack[lastExpr])
+		case "avg":
+			v.exprStack[lastExpr] = logicalplan.Avg(v.exprStack[lastExpr])
 		default:
 			return fmt.Errorf("unhandled aggregate function %s", expr.F)
 		}


Test patch:
diff --git a/aggregate_test.go b/aggregate_test.go
index bba0209..d5c996a 100644
--- a/aggregate_test.go
+++ b/aggregate_test.go
@@ -104,6 +104,11 @@ func TestAggregateInconsistentSchema(t *testing.T) {
 			alias:   "value_count",
 			expVals: []int64{2, 1},
 		},
+		{
+			fn:      logicalplan.Avg,
+			alias:   "value_avg",
+			expVals: []int64{2, 1},
+		},
 	} {
 		t.Run(testCase.alias, func(t *testing.T) {
 			var res arrow.Record
diff --git a/logictest/testdata/aggregate/aggregate b/logictest/testdata/aggregate/aggregate
index b8c780c..864f8be 100644
--- a/logictest/testdata/aggregate/aggregate
+++ b/logictest/testdata/aggregate/aggregate
@@ -28,6 +28,17 @@ select count(value) as value_count group by labels.label2
 ----
 value2  6
 
+exec
+select avg(value) as value_avg group by labels.label2
+----
+value2  3
+
+exec
+select avg(value) as value_avg group by labels.label4
+----
+        3
+value4  4
+
 exec
 select sum(value), count(value) group by stacktrace
 ----