Repository: polarsignals/frostdb
Pull request #765 (2024-03-12 10:40:10)

Add LIMIT support to FrostDB queries
There's probably lots of opportunities to improve and push down limits with a query optimizer. 
We will follow up with PRs to improve this.



Patch:
diff --git a/pqarrow/arrowutils/sort.go b/pqarrow/arrowutils/sort.go
index 63392da..a0e2ac2 100644
--- a/pqarrow/arrowutils/sort.go
+++ b/pqarrow/arrowutils/sort.go
@@ -104,11 +104,11 @@ func Take(ctx context.Context, r arrow.Record, indices *array.Int32) (arrow.Reco
 		col := r.Column(i)
 		if d, ok := col.(*array.Dictionary); ok {
 			g.Go(func() error {
-				return takeDictColumn(ctx, d, i, resArr, indices)
+				return TakeDictColumn(ctx, d, i, resArr, indices)
 			})
 		} else {
 			g.Go(func() error {
-				return takeColumn(ctx, col, i, resArr, indices)
+				return TakeColumn(ctx, col, i, resArr, indices)
 			})
 		}
 	}
@@ -119,7 +119,7 @@ func Take(ctx context.Context, r arrow.Record, indices *array.Int32) (arrow.Reco
 	return array.NewRecord(r.Schema(), resArr, int64(indices.Len())), nil
 }
 
-func takeColumn(ctx context.Context, a arrow.Array, idx int, arr []arrow.Array, indices *array.Int32) error {
+func TakeColumn(ctx context.Context, a arrow.Array, idx int, arr []arrow.Array, indices *array.Int32) error {
 	r, err := compute.TakeArray(ctx, a, indices)
 	if err != nil {
 		return err
@@ -128,7 +128,7 @@ func takeColumn(ctx context.Context, a arrow.Array, idx int, arr []arrow.Array,
 	return nil
 }
 
-func takeDictColumn(ctx context.Context, a *array.Dictionary, idx int, arr []arrow.Array, indices *array.Int32) error {
+func TakeDictColumn(ctx context.Context, a *array.Dictionary, idx int, arr []arrow.Array, indices *array.Int32) error {
 	var add func(b *array.BinaryDictionaryBuilder, idx int) error
 	switch e := a.Dictionary().(type) {
 	case *array.String:
diff --git a/query/engine.go b/query/engine.go
index 37a195a..3784d21 100644
--- a/query/engine.go
+++ b/query/engine.go
@@ -18,6 +18,7 @@ type Builder interface {
 	Filter(expr logicalplan.Expr) Builder
 	Distinct(expr ...logicalplan.Expr) Builder
 	Project(projections ...logicalplan.Expr) Builder
+	Limit(expr logicalplan.Expr) Builder
 	Execute(ctx context.Context, callback func(ctx context.Context, r arrow.Record) error) error
 	Explain(ctx context.Context) (string, error)
 }
@@ -131,6 +132,17 @@ func (b LocalQueryBuilder) Project(
 	}
 }
 
+func (b LocalQueryBuilder) Limit(
+	expr logicalplan.Expr,
+) Builder {
+	return LocalQueryBuilder{
+		pool:        b.pool,
+		tracer:      b.tracer,
+		planBuilder: b.planBuilder.Limit(expr),
+		execOpts:    b.execOpts,
+	}
+}
+
 func (b LocalQueryBuilder) Execute(ctx context.Context, callback func(ctx context.Context, r arrow.Record) error) error {
 	ctx, span := b.tracer.Start(ctx, "LocalQueryBuilder/Execute")
 	defer span.End()
diff --git a/query/logicalplan/builder.go b/query/logicalplan/builder.go
index a9dc38e..6d0d149 100644
--- a/query/logicalplan/builder.go
+++ b/query/logicalplan/builder.go
@@ -133,6 +133,22 @@ func (b Builder) Distinct(
 	}
 }
 
+func (b Builder) Limit(expr Expr) Builder {
+	if expr == nil {
+		return b
+	}
+
+	return Builder{
+		err: b.err,
+		plan: &LogicalPlan{
+			Input: b.plan,
+			Limit: &Limit{
+				Expr: expr,
+			},
+		},
+	}
+}
+
 func (b Builder) Aggregate(
 	aggExpr []*AggregationFunction,
 	groupExprs []Expr,
diff --git a/query/logicalplan/logicalplan.go b/query/logicalplan/logicalplan.go
index 6ab845f..23b5878 100644
--- a/query/logicalplan/logicalplan.go
+++ b/query/logicalplan/logicalplan.go
@@ -24,6 +24,7 @@ type LogicalPlan struct {
 	Distinct    *Distinct
 	Projection  *Projection
 	Aggregation *Aggregation
+	Limit       *Limit
 }
 
 // Callback is a function that is called throughout a chain of operators
@@ -405,3 +406,11 @@ type Aggregation struct {
 func (a *Aggregation) String() string {
 	return "Aggregation " + fmt.Sprint(a.AggExprs) + " Group: " + fmt.Sprint(a.GroupExprs)
 }
+
+type Limit struct {
+	Expr Expr
+}
+
+func (l *Limit) String() string {
+	return "Limit" + " Expr: " + fmt.Sprint(l.Expr)
+}
diff --git a/query/logicalplan/validate.go b/query/logicalplan/validate.go
index 7212d40..7b2f60a 100644
--- a/query/logicalplan/validate.go
+++ b/query/logicalplan/validate.go
@@ -119,6 +119,9 @@ func ValidateSingleFieldSet(plan *LogicalPlan) *PlanValidationError {
 	if plan.Aggregation != nil {
 		fieldsSet = append(fieldsSet, 5)
 	}
+	if plan.Limit != nil {
+		fieldsSet = append(fieldsSet, 6)
+	}
 
 	if len(fieldsSet) != 1 {
 		fieldsFound := make([]string, 0)
diff --git a/query/physicalplan/limit.go b/query/physicalplan/limit.go
new file mode 100644
index 0000000..249ee3d
--- /dev/null
+++ b/query/physicalplan/limit.go
@@ -0,0 +1,148 @@
+package physicalplan
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/apache/arrow/go/v15/arrow"
+	"github.com/apache/arrow/go/v15/arrow/array"
+	"github.com/apache/arrow/go/v15/arrow/compute"
+	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v15/arrow/scalar"
+	"go.opentelemetry.io/otel/trace"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/polarsignals/frostdb/pqarrow/arrowutils"
+	"github.com/polarsignals/frostdb/query/logicalplan"
+)
+
+type Limiter struct {
+	pool   memory.Allocator
+	tracer trace.Tracer
+	next   PhysicalPlan
+
+	count uint64
+}
+
+func Limit(pool memory.Allocator, tracer trace.Tracer, expr logicalplan.Expr) (*Limiter, error) {
+	literal, ok := expr.(*logicalplan.LiteralExpr)
+	if !ok {
+		return nil, fmt.Errorf("expected literal expression, got %T", expr)
+	}
+
+	var count uint64
+	switch v := literal.Value.(type) {
+	case *scalar.Uint64:
+		count = v.Value
+	case *scalar.Int64:
+		count = uint64(v.Value)
+	default:
+		return nil, fmt.Errorf("expected limit count type, got %T", v)
+	}
+
+	return &Limiter{
+		pool:   pool,
+		tracer: tracer,
+		count:  count,
+	}, nil
+}
+
+func (l *Limiter) SetNext(next PhysicalPlan) { l.next = next }
+
+func (l *Limiter) Finish(ctx context.Context) error { return l.next.Finish(ctx) }
+
+func (l *Limiter) Close() { l.next.Close() }
+
+func (l *Limiter) Draw() *Diagram {
+	var child *Diagram
+	if l.next != nil {
+		child = l.next.Draw()
+	}
+	details := fmt.Sprintf("Limit(%d)", l.count)
+	return &Diagram{Details: details, Child: child}
+}
+
+func (l *Limiter) Callback(ctx context.Context, r arrow.Record) error {
+	if r.NumRows() == 0 {
+		return l.next.Callback(ctx, r)
+	}
+	if l.count == 0 {
+		newRecord := array.NewRecord(r.Schema(), nil, 0)
+		return l.next.Callback(ctx, newRecord)
+	}
+
+	if uint64(r.NumRows()) <= l.count {
+		return l.next.Callback(ctx, r)
+	}
+
+	// TODO: We should figure out a way to avoid copying the record here.
+	// Maybe we can use a different approach to limit the record.
+
+	indicesBuilder := array.NewInt32Builder(l.pool)
+	defer indicesBuilder.Release()
+
+	for i := int32(0); i < int32(l.count); i++ {
+		indicesBuilder.Append(i)
+	}
+	indices := indicesBuilder.NewInt32Array()
+	defer indices.Release()
+
+	// compute.Take doesn't support dictionaries. Use take on r when r does not have
+	// dictionary column.
+	var hasDictionary bool
+	for i := 0; i < int(r.NumCols()); i++ {
+		if r.Column(i).DataType().ID() == arrow.DICTIONARY {
+			hasDictionary = true
+			break
+		}
+	}
+	if !hasDictionary {
+		res, err := compute.Take(
+			ctx,
+			compute.TakeOptions{BoundsCheck: true},
+			compute.NewDatumWithoutOwning(r),
+			compute.NewDatumWithoutOwning(indices),
+		)
+		if err != nil {
+			return err
+		}
+		r.Release()
+		return l.next.Callback(ctx, res.(*compute.RecordDatum).Value)
+	}
+	resArr := make([]arrow.Array, r.NumCols())
+
+	defer func() {
+		for _, a := range resArr {
+			if a != nil {
+				a.Release()
+			}
+		}
+	}()
+	var g errgroup.Group
+	for i := 0; i < int(r.NumCols()); i++ {
+		i := i
+		col := r.Column(i)
+		if d, ok := col.(*array.Dictionary); ok {
+			g.Go(func() error {
+				return arrowutils.TakeDictColumn(ctx, d, i, resArr, indices)
+			})
+		} else {
+			g.Go(func() error {
+				return arrowutils.TakeColumn(ctx, col, i, resArr, indices)
+			})
+		}
+	}
+	err := g.Wait()
+	if err != nil {
+		return err
+	}
+
+	if err := l.next.Callback(
+		ctx,
+		array.NewRecord(r.Schema(), resArr, int64(indices.Len())),
+	); err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/query/physicalplan/physicalplan.go b/query/physicalplan/physicalplan.go
index b91b4c1..d828568 100644
--- a/query/physicalplan/physicalplan.go
+++ b/query/physicalplan/physicalplan.go
@@ -379,6 +379,36 @@ func Build(
 				prev = prev[0:1]
 				prev[0] = d
 			}
+		case plan.Limit != nil:
+			var sync *Synchronizer
+			if len(prev) > 1 {
+				// These limit operators need to be synchronized.
+				sync = Synchronize(len(prev))
+			}
+			for i := 0; i < len(prev); i++ {
+				d, err := Limit(pool, tracer, plan.Limit.Expr)
+				if err != nil {
+					visitErr = err
+					return false
+				}
+				prev[i].SetNext(d)
+				prev[i] = d
+				if sync != nil {
+					d.SetNext(sync)
+				}
+			}
+			if sync != nil {
+				// Plan a limit operator to run a limit on all the
+				// synchronized limits.
+				d, err := Limit(pool, tracer, plan.Limit.Expr)
+				if err != nil {
+					visitErr = err
+					return false
+				}
+				sync.SetNext(d)
+				prev = prev[0:1]
+				prev[0] = d
+			}
 		case plan.Filter != nil:
 			// Create a filter for each previous plan.
 			// Can be multiple filters or just a single
diff --git a/sqlparse/visitor.go b/sqlparse/visitor.go
index ad780e4..9a2d1b7 100644
--- a/sqlparse/visitor.go
+++ b/sqlparse/visitor.go
@@ -134,6 +134,18 @@ func (v *astVisitor) Enter(n ast.Node) (nRes ast.Node, skipChildren bool) {
 			//
 			// SELECT sum(value)/count(value) value_avg, (timestamp/1000) as timestamp_bucket group by timestamp_bucket
 			v.builder = v.builder.Project(postProjections...)
+
+			// Finally we check if the result should be limited.
+			if expr.Limit != nil {
+				expr.Limit.Count.Accept(v)
+				lastExpr, _ := pop(v.exprStack)
+				v.builder = v.builder.Limit(lastExpr)
+			}
+		case expr.Limit != nil:
+			expr.Limit.Count.Accept(v)
+			lastExpr, newExprs := pop(v.exprStack)
+			v.builder = v.builder.Project(newExprs...)
+			v.builder = v.builder.Limit(lastExpr)
 		case expr.Distinct:
 			v.builder = v.builder.Project(v.exprStack...)
 			v.builder = v.builder.Distinct(v.exprStack...)


Test patch:
diff --git a/logictest/testdata/exec/aggregate/aggregate b/logictest/testdata/exec/aggregate/aggregate
index 6b4620d..a471ddc 100644
--- a/logictest/testdata/exec/aggregate/aggregate
+++ b/logictest/testdata/exec/aggregate/aggregate
@@ -108,3 +108,17 @@ value3  value2  null    value4  3
 value4  value2  null    null    4
 value5  value2  value3  null    5
 value6  value2  null    value4  6
+
+# Limit but there are fewer rows
+exec
+select sum(value) as value_sum, labels.label3 group by labels.label3 limit 3
+----
+14      null
+7       value3
+
+# Actually limit to 1 row
+exec
+select sum(value) as value_sum, labels.label3 group by labels.label3 limit 1
+----
+14      null
+
diff --git a/logictest/testdata/exec/filter/filter b/logictest/testdata/exec/filter/filter
index 65b218e..3be201d 100644
--- a/logictest/testdata/exec/filter/filter
+++ b/logictest/testdata/exec/filter/filter
@@ -29,6 +29,11 @@ select labels, stacktrace, timestamp, value where timestamp <= 2
 value1  value2  null    null    stack1  1       1
 value2  value2  value3  null    stack1  2       2
 
+exec
+select labels, stacktrace, timestamp, value where timestamp <= 2 limit 1
+----
+value1  value2  null    null    stack1  1       1
+
 exec
 select labels, stacktrace, timestamp, value where timestamp > 2
 ----
diff --git a/logictest/testdata/exec/limit/limit b/logictest/testdata/exec/limit/limit
new file mode 100644
index 0000000..7919403
--- /dev/null
+++ b/logictest/testdata/exec/limit/limit
@@ -0,0 +1,39 @@
+createtable schema=default
+----
+
+insert cols=(labels.label1, labels.label2, labels.label3, labels.label4, labels.label5)
+value1  value1  null    null    value1
+value2  value2  value3  null    value1
+value3  value1  null    value4  value1
+----
+
+exec
+select labels.label3 limit 0
+----
+
+exec
+select labels.label1 limit 4 # no need to limit
+----
+value1
+value2
+value3
+
+exec
+select labels.label2 limit 3 # no need to limit - exact 3 rows are the limit
+----
+value1
+value2
+value1
+
+exec
+select labels.label3 limit 2
+----
+null
+value3
+
+exec
+select labels.label1, labels.label2 limit 5
+----
+value1  value1
+value2  value2
+value3  value1
diff --git a/logictest/testdata/plan/aggregate/aggregate b/logictest/testdata/plan/aggregate/aggregate
index 4be8a6a..aad2bf6 100644
--- a/logictest/testdata/plan/aggregate/aggregate
+++ b/logictest/testdata/plan/aggregate/aggregate
@@ -82,3 +82,9 @@ exec
 explain select labels.label1, (max(value) - min(value)) / 2, sum(value) / count(value), (sum(value) / count(value)) * 2 group by labels.label1
 ----
 TableScan [concurrent] - Projection (labels.label1, value) - HashAggregate (max(value),min(value),sum(value),count(value),sum(value),count(value) by labels.label1) - Synchronizer - HashAggregate (max(value),min(value),sum(value),count(value),sum(value),count(value) by labels.label1) - Projection (labels.label1, max(value) - min(value) / 2, sum(value) / count(value), sum(value) / count(value) * 2)
+
+# Make sure that the limit happens after the aggregation is done.
+exec
+explain select sum(value) as value_sum group by (example_type, labels.label1) limit 2
+----
+TableScan [concurrent] - Projection (value, example_type, labels.label1) - HashAggregate (sum(value) by example_type,labels.label1) - Synchronizer - HashAggregate (sum(value) by example_type,labels.label1) - Projection (sum(value) as value_sum) - Limit(2)
diff --git a/logictest/testdata/plan/filter/filter b/logictest/testdata/plan/filter/filter
index 41d0d00..0e01035 100644
--- a/logictest/testdata/plan/filter/filter
+++ b/logictest/testdata/plan/filter/filter
@@ -15,3 +15,8 @@ exec
 explain select stacktrace, value where labels.label1 not like 'ue2' and stacktrace like 'ack'
 ----
 TableScan [concurrent] - PredicateFilter ((labels.label1 not contains ue2 AND stacktrace contains ack)) - Projection (stacktrace, value) - Synchronizer
+
+exec
+explain select stacktrace, value where stacktrace LIKE 'ack' limit 1
+----
+TableScan [concurrent] - PredicateFilter (stacktrace contains ack) - Projection (stacktrace, value) - Limit(1) - Synchronizer - Limit(1)
diff --git a/logictest/testdata/plan/limit/limit b/logictest/testdata/plan/limit/limit
new file mode 100644
index 0000000..39838ae
--- /dev/null
+++ b/logictest/testdata/plan/limit/limit
@@ -0,0 +1,12 @@
+createtable schema=default
+----
+
+exec
+explain select labels.label1
+----
+TableScan [concurrent] - Projection (labels.label1) - Synchronizer
+
+exec
+explain select labels.label1 limit 1
+----
+TableScan [concurrent] - Projection (labels.label1) - Limit(1) - Synchronizer - Limit(1)