Repository: polarsignals/frostdb
Pull request #751 (2024-02-21 14:57:12)

query: add global aggregation function pushdown optimization for Max
When reading Particulates from blocks (either in-memory or object storage), we can drastically reduce the runtime of global aggregations with no filters. This commit implements the optimization pushdown as a filter.

In case of Max (implemented in this commit) for example, the TrueNegativeFilter memoizes the maximum value seen and filters out any Particulates that will definitely not contribute to the aggregation function (i.e. whose max value is less than the max value seen so far). This drastically reduces global Max queries. For example, on 1.5GB of data, a global Max on an unsorted int64 column went from 900ms to 50ms.


Patch:
diff --git a/query/expr/filter.go b/query/expr/filter.go
index 817a1ae..65599ac 100644
--- a/query/expr/filter.go
+++ b/query/expr/filter.go
@@ -3,6 +3,9 @@ package expr
 import (
 	"errors"
 	"fmt"
+	"strings"
+	"sync"
+	"sync/atomic"
 
 	"github.com/parquet-go/parquet-go"
 
@@ -33,6 +36,7 @@ type Particulate interface {
 }
 
 type TrueNegativeFilter interface {
+	// Eval should be safe to call concurrently.
 	Eval(Particulate) (bool, error)
 }
 
@@ -54,7 +58,7 @@ func binaryBooleanExpr(expr *logicalplan.BinaryExpr) (TrueNegativeFilter, error)
 		fallthrough
 	case logicalplan.OpGtEq:
 		fallthrough
-	case logicalplan.OpEq: //, logicalplan.OpNotEq, logicalplan.OpLt, logicalplan.OpLtEq, logicalplan.OpGt, logicalplan.OpGtEq, logicalplan.OpRegexMatch, logicalplan.RegexNotMatch:
+	case logicalplan.OpEq: // , logicalplan.OpNotEq, logicalplan.OpLt, logicalplan.OpLtEq, logicalplan.OpGt, logicalplan.OpGtEq, logicalplan.OpRegexMatch, logicalplan.RegexNotMatch:
 		var leftColumnRef *ColumnRef
 		expr.Left.Accept(PreExprVisitorFunc(func(expr logicalplan.Expr) bool {
 			switch e := expr.(type) {
@@ -127,6 +131,80 @@ func binaryBooleanExpr(expr *logicalplan.BinaryExpr) (TrueNegativeFilter, error)
 	}
 }
 
+func aggregationExpr(expr *logicalplan.AggregationFunction) (TrueNegativeFilter, error) {
+	switch expr.Func {
+	case logicalplan.AggFuncMax:
+		a := &MaxAgg{}
+		expr.Expr.Accept(PreExprVisitorFunc(func(expr logicalplan.Expr) bool {
+			switch e := expr.(type) {
+			case *logicalplan.Column:
+				a.columnName = e.ColumnName
+			case *logicalplan.DynamicColumn:
+				a.columnName = e.ColumnName
+				a.dynamic = true
+			default:
+				return true
+			}
+			return false
+		}))
+		return a, nil
+	default:
+		return &AlwaysTrueFilter{}, nil
+	}
+}
+
+type MaxAgg struct {
+	maxMap sync.Map
+
+	// It would be nicer to use a dynamic-aware ColumnRef here, but that would
+	// introduce allocations (slices for indexes and concrete names), so for
+	// performance reasons we execute the column lookup manually.
+	columnName string
+	dynamic    bool
+}
+
+func (a *MaxAgg) Eval(p Particulate) (bool, error) {
+	processFurther := false
+	for i, f := range p.Schema().Fields() {
+		if (a.dynamic && !strings.HasPrefix(f.Name(), a.columnName+".")) || (!a.dynamic && f.Name() != a.columnName) {
+			continue
+		}
+
+		chunk := p.ColumnChunks()[i]
+		index, err := chunk.ColumnIndex()
+		if err != nil {
+			return false, fmt.Errorf("error retrieving column index in MaxAgg.Eval")
+		}
+		if NullCount(index) == chunk.NumValues() {
+			// This page is full of nulls. Nothing to do since we can't trust
+			// min/max index values. This chunk should not be processed since
+			// it can't contribute to min/max unless another column can.
+			continue
+		}
+
+		columnPointer, _ := a.maxMap.LoadOrStore(f.Name(), &atomic.Pointer[parquet.Value]{})
+		atomicMax := columnPointer.(*atomic.Pointer[parquet.Value])
+
+		v := Max(index)
+		for globalMax := atomicMax.Load(); globalMax == nil || compare(v, *globalMax) > 0; globalMax = atomicMax.Load() {
+			if atomicMax.CompareAndSwap(globalMax, &v) {
+				// At least one column exceeded the current max so this chunk
+				// satisfies the filter. Note that we do not break out of
+				// scanning the rest of the columns since we do want to memoize
+				// the max for other columns as well.
+				processFurther = true
+				break
+			}
+		}
+		if !a.dynamic && processFurther {
+			// No need to look at the remaining columns if we're only looking
+			// for a single concrete column.
+			break
+		}
+	}
+	return processFurther, nil
+}
+
 type AndExpr struct {
 	Left  TrueNegativeFilter
 	Right TrueNegativeFilter
@@ -178,6 +256,10 @@ func BooleanExpr(expr logicalplan.Expr) (TrueNegativeFilter, error) {
 	switch e := expr.(type) {
 	case *logicalplan.BinaryExpr:
 		return binaryBooleanExpr(e)
+	case *logicalplan.AggregationFunction:
+		// NOTE: Aggregations are optimized in the case of no grouping columns
+		// or other filters.
+		return aggregationExpr(e)
 	default:
 		return nil, fmt.Errorf("unsupported boolean expression %T", e)
 	}
diff --git a/query/logicalplan/optimize.go b/query/logicalplan/optimize.go
index f0a08c9..620eeae 100644
--- a/query/logicalplan/optimize.go
+++ b/query/logicalplan/optimize.go
@@ -15,6 +15,7 @@ func DefaultOptimizers() []Optimizer {
 		},
 		&FilterPushDown{},
 		&DistinctPushDown{},
+		&AggFuncPushDown{},
 	}
 }
 
@@ -154,3 +155,39 @@ func (p *DistinctPushDown) optimize(plan *LogicalPlan, distinctColumns []Expr) {
 		p.optimize(plan.Input, distinctColumns)
 	}
 }
+
+// AggFuncPushDown optimizer tries to push down an aggregation function operator
+// to the table provider. This can be done in the case of some aggregation
+// functions on global aggregations (i.e. no group by) without filters.
+// The storage engine can make smarter decisions than just returning all the
+// data, such as in the case of max functions, memoizing the max value seen
+// so far and only scanning row groups that contain a value greater than the
+// memoized value. It modifies the plan in place.
+type AggFuncPushDown struct{}
+
+func (p *AggFuncPushDown) Optimize(plan *LogicalPlan) *LogicalPlan {
+	p.optimize(plan, nil)
+	return plan
+}
+
+func (p *AggFuncPushDown) optimize(plan *LogicalPlan, filterExpr Expr) {
+	switch {
+	case plan.TableScan != nil:
+		if filterExpr != nil {
+			plan.TableScan.Filter = filterExpr
+		}
+	case plan.Aggregation != nil:
+		if len(plan.Aggregation.GroupExprs) == 0 && len(plan.Aggregation.AggExprs) == 1 {
+			// TODO(asubiotto): Should we make this less specific?
+			filterExpr = plan.Aggregation.AggExprs[0]
+		}
+	default:
+		// If we find anything other than a table scan after a global
+		// aggregation, bail out by setting the filterExpr to nil.
+		filterExpr = nil
+	}
+
+	if plan.Input != nil {
+		p.optimize(plan.Input, filterExpr)
+	}
+}


Test patch:
diff --git a/query/expr/filter_test.go b/query/expr/filter_test.go
new file mode 100644
index 0000000..4003b37
--- /dev/null
+++ b/query/expr/filter_test.go
@@ -0,0 +1,209 @@
+package expr
+
+import (
+	"testing"
+
+	"github.com/parquet-go/parquet-go"
+	"github.com/stretchr/testify/require"
+)
+
+type fakeParticulate struct {
+	schema       *parquet.Schema
+	columnChunks []parquet.ColumnChunk
+}
+
+// newFakeParticulate creates a new fake particulate. maxValues[i] is the
+// maximum value of the ith column. If maxValues[i] is -1, then the ith column
+// index will return that it is full of nulls.
+func newFakeParticulate(columnNames []string, maxValues []int64) fakeParticulate {
+	if len(columnNames) != len(maxValues) {
+		panic("columnNames and maxValues must have the same length")
+	}
+	vForName := make(map[string]int64)
+	for i, name := range columnNames {
+		vForName[name] = maxValues[i]
+	}
+
+	g := parquet.Group{}
+	for _, name := range columnNames {
+		g[name] = parquet.Int(64)
+	}
+	s := parquet.NewSchema("", g)
+
+	columnChunks := make([]parquet.ColumnChunk, len(columnNames))
+	// Iterate over the schema in creation order (not doing so mixes up which
+	// column has which value).
+	for i, f := range s.Fields() {
+		numValues := int64(1)
+		numNulls := int64(0)
+		var maxV parquet.Value
+		if v := vForName[f.Name()]; v == -1 {
+			maxV = parquet.ValueOf(nil)
+			numNulls = numValues
+		} else {
+			maxV = parquet.ValueOf(v)
+		}
+		columnChunks[i] = &FakeColumnChunk{
+			index: &FakeColumnIndex{
+				numPages:  1,
+				min:       parquet.Value{},
+				max:       maxV,
+				nullCount: numNulls,
+			},
+			numValues: numValues,
+		}
+	}
+	return fakeParticulate{
+		schema:       s,
+		columnChunks: columnChunks,
+	}
+}
+
+func (f fakeParticulate) Schema() *parquet.Schema {
+	return f.schema
+}
+
+func (f fakeParticulate) ColumnChunks() []parquet.ColumnChunk {
+	return f.columnChunks
+}
+
+func TestMaxAgg(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		agg  *MaxAgg
+		// values[i] is the ith particulate.
+		values []struct {
+			p fakeParticulate
+			// expected result of Eval(values[i]).
+			expected bool
+		}
+	}{
+		{
+			name: "ConcreteColumn",
+			agg: &MaxAgg{
+				columnName: "a",
+			},
+			values: []struct {
+				p        fakeParticulate
+				expected bool
+			}{
+				{
+					newFakeParticulate(
+						[]string{"anotfound"},
+						[]int64{3},
+					),
+					// Column not found.
+					false,
+				},
+				{
+					newFakeParticulate(
+						[]string{"a.a"},
+						[]int64{3},
+					),
+					// Column not found.
+					false,
+				},
+				{
+					newFakeParticulate(
+						[]string{"a"},
+						[]int64{2},
+					),
+					true,
+				},
+				{
+					newFakeParticulate(
+						[]string{"a"},
+						[]int64{1},
+					),
+					// Values is less than the current max.
+					false,
+				},
+				{
+					newFakeParticulate(
+						[]string{"b.a", "a"},
+						[]int64{-1, 3},
+					),
+					true,
+				},
+			},
+		},
+		{
+			name: "DynamicColumn",
+			agg: &MaxAgg{
+				columnName: "a",
+				dynamic:    true,
+			},
+			values: []struct {
+				p        fakeParticulate
+				expected bool
+			}{
+				{
+					newFakeParticulate(
+						[]string{"a"},
+						[]int64{3},
+					),
+					// Not a dynamic column.
+					false,
+				},
+				{
+					newFakeParticulate(
+						[]string{"a.a", "a"},
+						[]int64{-1, 3},
+					),
+					// Nulls in the dynamic column.
+					false,
+				},
+				{
+					newFakeParticulate(
+						[]string{"a.a", "a.b", "a.c"},
+						[]int64{3, 3, 3},
+					),
+					// This particulate should pass the filter. Also, this
+					// particulate verifies that the filter memoizes the max for
+					// all columns in the particulate for the future.
+					true,
+				},
+				{
+					newFakeParticulate(
+						[]string{"a.a"},
+						[]int64{1},
+					),
+					// Lower than max.
+					false,
+				},
+				{
+					newFakeParticulate(
+						[]string{"a.b"},
+						[]int64{1},
+					),
+					// Lower than max.
+					false,
+				},
+				{
+					newFakeParticulate(
+						[]string{"a.c"},
+						[]int64{1},
+					),
+					// Lower than max.
+					false,
+				},
+				{
+					newFakeParticulate(
+						[]string{"a.a", "a.b"},
+						[]int64{1, 4},
+					),
+					// First column is lower, but second column is higher.
+					true,
+				},
+			},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			for i, v := range tc.values {
+				got, err := tc.agg.Eval(v.p)
+				require.NoError(t, err)
+				require.Equal(t, v.expected, got, "index %d mismatch", i)
+			}
+		})
+	}
+}
diff --git a/query/logicalplan/optimize_test.go b/query/logicalplan/optimize_test.go
index cb9149d..838a02c 100644
--- a/query/logicalplan/optimize_test.go
+++ b/query/logicalplan/optimize_test.go
@@ -194,3 +194,77 @@ func TestAllOptimizers(t *testing.T) {
 		p.Input.Input.Input.TableScan,
 	)
 }
+
+func TestAggFuncPushDown(t *testing.T) {
+	t.Run("GlobalAgg", func(t *testing.T) {
+		tableProvider := &mockTableProvider{schema: dynparquet.NewSampleSchema()}
+		p, err := (&Builder{}).
+			Scan(tableProvider, "table1").
+			Aggregate(
+				[]*AggregationFunction{Max(Col("value"))},
+				nil,
+			).
+			Build()
+		require.NoError(t, err)
+
+		p = (&AggFuncPushDown{}).Optimize(p)
+		// Aggregation should still happen.
+		require.NotNil(t, p.Aggregation)
+		require.Equal(t,
+			&TableScan{
+				TableName:     "table1",
+				TableProvider: tableProvider,
+				Filter: &AggregationFunction{
+					Func: AggFuncMax,
+					Expr: &Column{ColumnName: "value"},
+				},
+			},
+			p.Input.TableScan,
+		)
+	})
+	t.Run("DontPushWithGroupExprs", func(t *testing.T) {
+		tableProvider := &mockTableProvider{schema: dynparquet.NewSampleSchema()}
+		p, err := (&Builder{}).
+			Scan(tableProvider, "table1").
+			Aggregate(
+				[]*AggregationFunction{Max(Col("value"))},
+				[]Expr{Col("stacktrace")},
+			).
+			Build()
+		require.NoError(t, err)
+
+		p = (&AggFuncPushDown{}).Optimize(p)
+		// Aggregation should still happen.
+		require.NotNil(t, p.Aggregation)
+		require.Equal(t,
+			&TableScan{
+				TableName:     "table1",
+				TableProvider: tableProvider,
+			},
+			p.Input.TableScan,
+		)
+	})
+	t.Run("DontPushWithFilter", func(t *testing.T) {
+		tableProvider := &mockTableProvider{schema: dynparquet.NewSampleSchema()}
+		p, err := (&Builder{}).
+			Scan(tableProvider, "table1").
+			Filter(Col("labels.test").Eq(Literal("abc"))).
+			Aggregate(
+				[]*AggregationFunction{Max(Col("value"))},
+				nil,
+			).
+			Build()
+		require.NoError(t, err)
+
+		p = (&AggFuncPushDown{}).Optimize(p)
+		// Aggregation should still happen.
+		require.NotNil(t, p.Aggregation)
+		require.Equal(t,
+			&TableScan{
+				TableName:     "table1",
+				TableProvider: tableProvider,
+			},
+			p.Input.Input.TableScan,
+		)
+	})
+}