Repository: polarsignals/frostdb
Pull request #501 (2023-08-25 21:39:29)

Aggregation dyn cols
Adds support for performing aggregations on dynamic columns.


Patch:
diff --git a/query/logicalplan/validate.go b/query/logicalplan/validate.go
index 237594a..83e3a67 100644
--- a/query/logicalplan/validate.go
+++ b/query/logicalplan/validate.go
@@ -166,6 +166,10 @@ func ValidateAggregation(plan *LogicalPlan) *PlanValidationError {
 	return nil
 }
 
+type Named interface {
+	Name() string
+}
+
 func ValidateAggregationExpr(plan *LogicalPlan) *ExprValidationError {
 	aliases := map[string]struct{}{}
 
@@ -174,10 +178,13 @@ func ValidateAggregationExpr(plan *LogicalPlan) *ExprValidationError {
 		colFinder := newTypeFinder((*Column)(nil))
 		expr.Accept(&colFinder)
 
+		dynColFinder := newTypeFinder((*DynamicColumn)(nil))
+		expr.Accept(&dynColFinder)
+
 		aggFuncFinder := newTypeFinder((*AggregationFunction)(nil))
 		expr.Accept(&aggFuncFinder)
 
-		if colFinder.result == nil || aggFuncFinder.result == nil {
+		if (colFinder.result == nil && dynColFinder.result == nil) || aggFuncFinder.result == nil {
 			return &ExprValidationError{
 				message: "aggregation expression is invalid. must contain AggregationFunction and Column",
 				expr:    expr,
@@ -185,16 +192,21 @@ func ValidateAggregationExpr(plan *LogicalPlan) *ExprValidationError {
 		}
 
 		// check that column being aggregated on exists in the schema
-		colExpr := colFinder.result.(*Column)
 		schema := plan.InputSchema()
 		if schema == nil {
 			return nil // cannot check column type if there's no input schema
 		}
 
-		column, found := schema.ColumnByName(colExpr.ColumnName)
+		var named Named
+		named = colFinder.result
+		if named == nil {
+			named = dynColFinder.result
+		}
+
+		column, found := schema.ColumnByName(named.Name())
 		if !found {
 			return &ExprValidationError{
-				message: fmt.Sprintf("column not found: %s", colExpr.ColumnName),
+				message: fmt.Sprintf("column not found: %s", named.Name()),
 				expr:    expr,
 			}
 		}
diff --git a/query/physicalplan/aggregate.go b/query/physicalplan/aggregate.go
index b054743..635f9fe 100644
--- a/query/physicalplan/aggregate.go
+++ b/query/physicalplan/aggregate.go
@@ -46,6 +46,10 @@ func Aggregate(
 			case *logicalplan.Column:
 				aggregation.expr = e
 				aggColumnFound = true
+			case *logicalplan.DynamicColumn:
+				aggregation.expr = e
+				aggregation.dynamic = true
+				aggColumnFound = true
 			}
 
 			return true
@@ -127,6 +131,7 @@ func chooseAggregationFunction(
 // Aggregation groups together some lower level primitives to for the column to be aggregated by its function.
 type Aggregation struct {
 	expr       logicalplan.Expr
+	dynamic    bool // dynamic indicates that this aggregation is performed against a dynamic column.
 	resultName string
 	function   logicalplan.AggFunc
 	arrays     []builder.ColumnBuilder // TODO: These can actually live outside this struct and be shared. Only at the very end will they be read by each column and then aggregated separately.
@@ -164,8 +169,9 @@ type hashtuple struct {
 
 // hashAggregate represents a single hash aggregation.
 type hashAggregate struct {
-	aggregations []Aggregation
-	groupByCols  map[string]builder.ColumnBuilder
+	dynamicAggregations []Aggregation
+	aggregations        []Aggregation
+	groupByCols         map[string]builder.ColumnBuilder
 
 	colOrdering []string
 	rowCount    int
@@ -178,6 +184,16 @@ func NewHashAggregate(
 	groupByColumnMatchers []logicalplan.Expr,
 	finalStage bool,
 ) *HashAggregate {
+	dynamic := []Aggregation{}
+	static := []Aggregation{}
+	for _, agg := range aggregations {
+		if agg.dynamic {
+			dynamic = append(dynamic, agg)
+		} else {
+			static = append(static, agg)
+		}
+	}
+
 	return &HashAggregate{
 		pool:   pool,
 		tracer: tracer,
@@ -192,9 +208,10 @@ func NewHashAggregate(
 		hashToAggregate:    map[uint64]hashtuple{},
 		aggregates: []*hashAggregate{ // initialize a single hash aggregate; we expect this array to only every grow during very large aggregations.
 			{
-				aggregations: aggregations,
-				groupByCols:  map[string]builder.ColumnBuilder{},
-				colOrdering:  []string{},
+				dynamicAggregations: dynamic,
+				aggregations:        static,
+				groupByCols:         map[string]builder.ColumnBuilder{},
+				colOrdering:         []string{},
 			},
 		},
 	}
@@ -421,11 +438,38 @@ func (a *HashAggregate) Callback(ctx context.Context, r arrow.Record) error {
 			}
 		}
 
+		for _, col := range aggregate.dynamicAggregations {
+			if a.finalStage {
+				if col.expr.MatchColumn(field.Name) {
+					// expand the aggregate.aggregations with a final concrete column aggregation.
+					columnToAggregate = append(columnToAggregate, nil)
+					aggregate.aggregations = append(aggregate.aggregations, Aggregation{
+						expr:       logicalplan.Col(field.Name),
+						dynamic:    true,
+						resultName: resultNameWithConcreteColumn(col.function, field.Name),
+						function:   col.function,
+					})
+				}
+			} else {
+				// If we're aggregating the raw data we need to find the columns by their actual names for now.
+				if col.expr.MatchColumn(field.Name) {
+					// expand the aggregate.aggregations with a concrete column aggregation.
+					columnToAggregate = append(columnToAggregate, nil)
+					aggregate.aggregations = append(aggregate.aggregations, Aggregation{
+						expr:       logicalplan.Col(field.Name),
+						dynamic:    true,
+						resultName: field.Name, // Don't rename the column yet, we'll do that in the final stage. Dynamic aggregations can't match agains't the pre-computed name.
+						function:   col.function,
+					})
+				}
+			}
+		}
+
 		for j, col := range aggregate.aggregations {
 			// If we're aggregating at the final stage we have previously
 			// renamed the pre-aggregated columns to their result names.
 			if a.finalStage {
-				if col.resultName == field.Name {
+				if col.resultName == field.Name || (col.dynamic && col.expr.MatchColumn(field.Name)) {
 					columnToAggregate[j] = r.Column(i)
 					aggregateFieldsFound++
 				}
@@ -815,3 +859,20 @@ func runAggregation(finalStage bool, fn logicalplan.AggFunc, pool memory.Allocat
 	}
 	return aggFunc.Aggregate(pool, arrs)
 }
+
+func resultNameWithConcreteColumn(function logicalplan.AggFunc, col string) string {
+	switch function {
+	case logicalplan.AggFuncSum:
+		return logicalplan.Sum(logicalplan.Col(col)).Name()
+	case logicalplan.AggFuncMin:
+		return logicalplan.Min(logicalplan.Col(col)).Name()
+	case logicalplan.AggFuncMax:
+		return logicalplan.Max(logicalplan.Col(col)).Name()
+	case logicalplan.AggFuncCount:
+		return logicalplan.Count(logicalplan.Col(col)).Name()
+	case logicalplan.AggFuncAvg:
+		return logicalplan.Avg(logicalplan.Col(col)).Name()
+	default:
+		return ""
+	}
+}


Test patch:
diff --git a/aggregate_test.go b/aggregate_test.go
index f2f0480..9400e0c 100644
--- a/aggregate_test.go
+++ b/aggregate_test.go
@@ -2,6 +2,7 @@ package frostdb
 
 import (
 	"context"
+	"math/rand"
 	"sort"
 	"strconv"
 	"testing"
@@ -13,6 +14,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/polarsignals/frostdb/dynparquet"
+	schemapb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/schema/v1alpha1"
 	"github.com/polarsignals/frostdb/query"
 	"github.com/polarsignals/frostdb/query/logicalplan"
 )
@@ -335,3 +337,75 @@ func BenchmarkAggregation(b *testing.B) {
 		})
 	}
 }
+
+func Test_Aggregation_DynCol(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+	defer mem.AssertSize(t, 0)
+
+	config := NewTableConfig(
+		&schemapb.Schema{
+			Name: "test",
+			Columns: []*schemapb.Column{{
+				Name: "foo",
+				StorageLayout: &schemapb.StorageLayout{
+					Type:     schemapb.StorageLayout_TYPE_INT64,
+					Nullable: true,
+				},
+				Dynamic: false,
+			}},
+			SortingColumns: []*schemapb.SortingColumn{{
+				Name:      "foo",
+				Direction: schemapb.SortingColumn_DIRECTION_ASCENDING,
+			}},
+		},
+	)
+
+	logger := newTestLogger(t)
+
+	c, err := New(
+		WithLogger(logger),
+	)
+	require.NoError(t, err)
+	defer c.Close()
+	db, err := c.DB(context.Background(), "test")
+	require.NoError(t, err)
+	table, err := db.Table("test", config)
+	require.NoError(t, err)
+
+	schema := arrow.NewSchema([]arrow.Field{
+		{Name: "foo.bar", Type: arrow.PrimitiveTypes.Int64},
+		{Name: "foo.baz", Type: arrow.PrimitiveTypes.Int64},
+		{Name: "foo.bah", Type: arrow.PrimitiveTypes.Int64},
+	}, nil)
+
+	bldr := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+	defer bldr.Release()
+
+	n := 3
+	for i := 0; i < schema.NumFields(); i++ {
+		for j := 0; j < n; j++ {
+			bldr.Field(i).(*array.Int64Builder).Append(int64(rand.Intn(100)))
+		}
+	}
+
+	r := bldr.NewRecord()
+	defer r.Release()
+
+	ctx := context.Background()
+	_, err = table.InsertRecord(ctx, r)
+	require.NoError(t, err)
+
+	engine := query.NewEngine(mem, db.TableProvider())
+
+	err = engine.ScanTable("test").
+		Aggregate(
+			[]logicalplan.Expr{logicalplan.Max(logicalplan.DynCol("foo"))},
+			nil,
+		).
+		Execute(context.Background(), func(ctx context.Context, ar arrow.Record) error {
+			require.Equal(t, 3, int(ar.NumCols()))
+			require.Equal(t, 1, int(ar.NumRows()))
+			return nil
+		})
+	require.NoError(t, err)
+}