Add experimental pushdown for dot notation (#4049)

FerretDB · Feb 12, 2024 · e92ded7 · e92ded7
1 parent 282c8e1
commit e92ded7
Show file tree

Hide file tree

Showing 11 changed files with 81 additions and 19 deletions.
diff --git a/cmd/ferretdb/main.go b/cmd/ferretdb/main.go
@@ -92,7 +92,8 @@ var cli struct {
 	Test struct {
 		RecordsDir string `default:"" help:"Testing: directory for record files."`
 
-		DisablePushdown bool `default:"false" help:"Experimental: disable pushdown."`
+		DisablePushdown      bool `default:"false" help:"Experimental: disable pushdown."`
+		EnableNestedPushdown bool `default:"false" help:"Experimental: enable pushdown for dot notation."`
 
 		CappedCleanup struct {
 			Interval   time.Duration `default:"1m" help:"Experimental: capped collections cleanup interval."`
@@ -353,6 +354,10 @@ func run() {
 
 	var wg sync.WaitGroup
 
+	if cli.Test.DisablePushdown && cli.Test.EnableNestedPushdown {
+		logger.Sugar().Fatal("--test-disable-pushdown and --test-enable-nested-pushdown should not be set at the same time")
+	}
+
 	// https://github.com/alecthomas/kong/issues/389
 	if cli.DebugAddr != "" && cli.DebugAddr != "-" {
 		wg.Add(1)
@@ -403,6 +408,7 @@ func run() {
 
 		TestOpts: registry.TestOpts{
 			DisablePushdown:         cli.Test.DisablePushdown,
+			EnableNestedPushdown:    cli.Test.EnableNestedPushdown,
 			CappedCleanupInterval:   cli.Test.CappedCleanup.Interval,
 			CappedCleanupPercentage: cli.Test.CappedCleanup.Percentage,
 			EnableNewAuth:           cli.Test.EnableNewAuth,

diff --git a/internal/backends/postgresql/query.go b/internal/backends/postgresql/query.go
@@ -103,6 +103,13 @@ func prepareWhereClause(p *metadata.Placeholder, sqlFilters *types.Document) (st
 			return "", nil, lazyerrors.Error(err)
 		}
 
+		keyOperator := "->" // keyOperator is the operator that is used to access the field. (->/#>)
+
+		// key can be either a string '"v"' or PostgreSQL path '{v,foo}'.
+		// We use path type only for dot notation due to simplicity of SQL queries, and the fact
+		// that path doesn't handle empty keys.
+		var key any = rootKey
+
 		// don't pushdown $comment, as it's attached to query with select clause
 		//
 		// all of the other top-level operators such as `$or` do not support pushdown yet
@@ -116,11 +123,11 @@ func prepareWhereClause(p *metadata.Placeholder, sqlFilters *types.Document) (st
 
 		switch {
 		case err == nil:
-			// Handle dot notation.
-			// TODO https://github.com/FerretDB/FerretDB/issues/2069
 			if path.Len() > 1 {
-				continue
+				keyOperator = "#>"
+				key = path.Slice() // '{v,foo}'
 			}
+
 		case errors.As(err, &pe):
 			// ignore empty key error, otherwise return error
 			if pe.Code() != types.ErrPathElementEmpty {
@@ -148,7 +155,7 @@ func prepareWhereClause(p *metadata.Placeholder, sqlFilters *types.Document) (st
 
 				switch k {
 				case "$eq":
-					if f, a := filterEqual(p, rootKey, v); f != "" {
+					if f, a := filterEqual(p, key, v, keyOperator); f != "" {
 						filters = append(filters, f)
 						args = append(args, a...)
 					}
@@ -209,7 +216,7 @@ func prepareWhereClause(p *metadata.Placeholder, sqlFilters *types.Document) (st
 			// type not supported for pushdown
 
 		case float64, string, types.ObjectID, bool, time.Time, int32, int64:
-			if f, a := filterEqual(p, rootKey, v); f != "" {
+			if f, a := filterEqual(p, key, v, keyOperator); f != "" {
 				filters = append(filters, f)
 				args = append(args, a...)
 			}
@@ -253,9 +260,9 @@ func prepareOrderByClause(p *metadata.Placeholder, sort *types.Document) (string
 
 // filterEqual returns the proper SQL filter with arguments that filters documents
 // where the value under k is equal to v.
-func filterEqual(p *metadata.Placeholder, k string, v any) (filter string, args []any) {
+func filterEqual(p *metadata.Placeholder, k any, v any, operator string) (filter string, args []any) {
 	// Select if value under the key is equal to provided value.
-	sql := `%[1]s->%[2]s @> %[3]s`
+	sql := `%[1]s%[2]s%[3]s @> %[4]s`
 
 	switch v := v.(type) {
 	case *types.Document, *types.Array, types.Binary,
@@ -267,33 +274,33 @@ func filterEqual(p *metadata.Placeholder, k string, v any) (filter string, args
 		// TODO https://github.com/FerretDB/FerretDB/issues/3626
 		switch {
 		case v > types.MaxSafeDouble:
-			sql = `%[1]s->%[2]s > %[3]s`
+			sql = `%[1]s%[2]s%[3]s > %[4]s`
 			v = types.MaxSafeDouble
 
 		case v < -types.MaxSafeDouble:
-			sql = `%[1]s->%[2]s < %[3]s`
+			sql = `%[1]s%[2]s%[3]s < %[4]s`
 			v = -types.MaxSafeDouble
 		default:
 			// don't change the default eq query
 		}
 
-		filter = fmt.Sprintf(sql, metadata.DefaultColumn, p.Next(), p.Next())
+		filter = fmt.Sprintf(sql, metadata.DefaultColumn, operator, p.Next(), p.Next())
 		args = append(args, k, v)
 
 	case string, types.ObjectID, time.Time:
 		// merge with the case below?
 		// TODO https://github.com/FerretDB/FerretDB/issues/3626
 
 		// don't change the default eq query
-		filter = fmt.Sprintf(sql, metadata.DefaultColumn, p.Next(), p.Next())
+		filter = fmt.Sprintf(sql, metadata.DefaultColumn, operator, p.Next(), p.Next())
 		args = append(args, k, string(must.NotFail(sjson.MarshalSingleValue(v))))
 
 	case bool, int32:
 		// merge with the case above?
 		// TODO https://github.com/FerretDB/FerretDB/issues/3626
 
 		// don't change the default eq query
-		filter = fmt.Sprintf(sql, metadata.DefaultColumn, p.Next(), p.Next())
+		filter = fmt.Sprintf(sql, metadata.DefaultColumn, operator, p.Next(), p.Next())
 		args = append(args, k, v)
 
 	case int64:
@@ -303,17 +310,17 @@ func filterEqual(p *metadata.Placeholder, k string, v any) (filter string, args
 		// If value cannot be safe double, fetch all numbers out of the safe range.
 		switch {
 		case v > maxSafeDouble:
-			sql = `%[1]s->%[2]s > %[3]s`
+			sql = `%[1]s%[2]s%[3]s > %[4]s`
 			v = maxSafeDouble
 
 		case v < -maxSafeDouble:
-			sql = `%[1]s->%[2]s < %[3]s`
+			sql = `%[1]s%[2]s%[3]s < %[4]s`
 			v = -maxSafeDouble
 		default:
 			// don't change the default eq query
 		}
 
-		filter = fmt.Sprintf(sql, metadata.DefaultColumn, p.Next(), p.Next())
+		filter = fmt.Sprintf(sql, metadata.DefaultColumn, operator, p.Next(), p.Next())
 		args = append(args, k, v)
 
 	default:

diff --git a/internal/backends/postgresql/query_test.go b/internal/backends/postgresql/query_test.go
@@ -95,6 +95,8 @@ func TestPrepareWhereClause(t *testing.T) {
 
 	// WHERE clauses occurring frequently in tests
 	whereContain := " WHERE _jsonb->$1 @> $2"
+	whereContainDotNotation := " WHERE _jsonb#>$1 @> $2"
+
 	whereGt := " WHERE _jsonb->$1 > $2"
 	whereNotEq := ` WHERE NOT ( _jsonb ? $1 AND _jsonb->$1 @> $2 AND _jsonb->'$s'->'p'->$1->'t' = `
 
@@ -117,14 +119,17 @@ func TestPrepareWhereClause(t *testing.T) {
 			expected: whereContain,
 		},
 		"IDDotNotation": {
-			filter: must.NotFail(types.NewDocument("_id.doc", "foo")),
+			filter:   must.NotFail(types.NewDocument("_id.doc", "foo")),
+			expected: whereContainDotNotation,
 		},
 
 		"DotNotation": {
-			filter: must.NotFail(types.NewDocument("v.doc", "foo")),
+			filter:   must.NotFail(types.NewDocument("v.doc", "foo")),
+			expected: whereContainDotNotation,
 		},
 		"DotNotationArrayIndex": {
-			filter: must.NotFail(types.NewDocument("v.arr.0", "foo")),
+			filter:   must.NotFail(types.NewDocument("v.arr.0", "foo")),
+			expected: whereContainDotNotation,
 		},
 
 		"ImplicitString": {

diff --git a/internal/handler/handler.go b/internal/handler/handler.go
@@ -77,6 +77,7 @@ type NewOpts struct {
 
 	// test options
 	DisablePushdown         bool
+	EnableNestedPushdown    bool
 	CappedCleanupInterval   time.Duration
 	CappedCleanupPercentage uint8
 	EnableNewAuth           bool

diff --git a/internal/handler/msg_aggregate.go b/internal/handler/msg_aggregate.go
@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"math"
 	"os"
+	"strings"
 	"time"
 
 	"go.uber.org/zap"
@@ -281,6 +282,18 @@ func (h *Handler) MsgAggregate(ctx context.Context, msg *wire.OpMsg) (*wire.OpMs
 			qp.Filter = filter
 		}
 
+		if !h.EnableNestedPushdown && filter != nil {
+			qp.Filter = filter.DeepCopy()
+
+			for _, k := range qp.Filter.Keys() {
+				if !strings.ContainsRune(k, '.') {
+					continue
+				}
+
+				qp.Filter.Remove(k)
+			}
+		}
+
 		if sort, err = common.ValidateSortDocument(sort); err != nil {
 			closer.Close()
 

diff --git a/internal/handler/msg_explain.go b/internal/handler/msg_explain.go
@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/FerretDB/FerretDB/build/version"
 	"github.com/FerretDB/FerretDB/internal/backends"
@@ -90,6 +91,18 @@ func (h *Handler) MsgExplain(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg,
 		qp.Filter = params.Filter
 	}
 
+	if !h.EnableNestedPushdown && params.Filter != nil {
+		qp.Filter = params.Filter.DeepCopy()
+
+		for _, k := range qp.Filter.Keys() {
+			if !strings.ContainsRune(k, '.') {
+				continue
+			}
+
+			qp.Filter.Remove(k)
+		}
+	}
+
 	if params.Sort, err = common.ValidateSortDocument(params.Sort); err != nil {
 		var pathErr *types.PathError
 		if errors.As(err, &pathErr) && pathErr.Code() == types.ErrPathElementEmpty {

diff --git a/internal/handler/msg_find.go b/internal/handler/msg_find.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strings"
 	"time"
 
 	"go.uber.org/zap"
@@ -221,6 +222,18 @@ func (h *Handler) makeFindQueryParams(params *common.FindParams, cInfo *backends
 		qp.Filter = params.Filter
 	}
 
+	if !h.EnableNestedPushdown && params.Filter != nil {
+		qp.Filter = params.Filter.DeepCopy()
+
+		for _, k := range qp.Filter.Keys() {
+			if !strings.ContainsRune(k, '.') {
+				continue
+			}
+
+			qp.Filter.Remove(k)
+		}
+	}
+
 	if params.Sort, err = common.ValidateSortDocument(params.Sort); err != nil {
 		var pathErr *types.PathError
 		if errors.As(err, &pathErr) && pathErr.Code() == types.ErrPathElementEmpty {

diff --git a/internal/handler/registry/mysql.go b/internal/handler/registry/mysql.go
@@ -41,6 +41,7 @@ func init() {
 			StateProvider: opts.StateProvider,
 
 			DisablePushdown:         opts.DisablePushdown,
+			EnableNestedPushdown:    opts.EnableNestedPushdown,
 			CappedCleanupPercentage: opts.CappedCleanupPercentage,
 			CappedCleanupInterval:   opts.CappedCleanupInterval,
 			EnableNewAuth:           opts.EnableNewAuth,

diff --git a/internal/handler/registry/postgresql.go b/internal/handler/registry/postgresql.go
@@ -43,6 +43,7 @@ func init() {
 			StateProvider: opts.StateProvider,
 
 			DisablePushdown:         opts.DisablePushdown,
+			EnableNestedPushdown:    opts.EnableNestedPushdown,
 			CappedCleanupPercentage: opts.CappedCleanupPercentage,
 			CappedCleanupInterval:   opts.CappedCleanupInterval,
 			EnableNewAuth:           opts.EnableNewAuth,

diff --git a/internal/handler/registry/registry.go b/internal/handler/registry/registry.go
@@ -67,6 +67,7 @@ type NewHandlerOpts struct {
 // TestOpts represents experimental configuration options.
 type TestOpts struct {
 	DisablePushdown         bool
+	EnableNestedPushdown    bool
 	CappedCleanupInterval   time.Duration
 	CappedCleanupPercentage uint8
 	EnableNewAuth           bool

diff --git a/internal/handler/registry/sqlite.go b/internal/handler/registry/sqlite.go
@@ -43,6 +43,7 @@ func init() {
 			StateProvider: opts.StateProvider,
 
 			DisablePushdown:         opts.DisablePushdown,
+			EnableNestedPushdown:    opts.EnableNestedPushdown,
 			CappedCleanupPercentage: opts.CappedCleanupPercentage,
 			CappedCleanupInterval:   opts.CappedCleanupInterval,
 			EnableNewAuth:           opts.EnableNewAuth,