Decode only fields observed by projections

Author: xet7

CHANGELOG.md

@@ -6,6 +6,15 @@
 
 ### Fixed 🐛
 
+- **Inclusion projections no longer decode fields the query cannot observe.**
+  SQLite still parses and filters the unchanged stored SJSON document, but its
+  iterator now recursively decodes only projected top-level fields plus fields
+  required by filters and sorts. Small projections such as `_id` no longer
+  decode large member arrays from every matching document during collection
+  scans. Exclusion projections retain full decoding. Unit tests cover nested
+  query-field collection, missing fields, stored field order and skipped nested
+  values by @xet7. Thanks to xet7.
+
 - **Tailable OpLog cursors wake on writes instead of polling SQLite.** The
   OpLog decorator now broadcasts each successful append to every `awaitData`
   waiter. Cursors retain that notification generation and wait before their

internal/backends/collection.go

@@ -78,6 +78,10 @@ type QueryParams struct {
 	Filter *types.Document
 	Sort   *types.Document
 	Limit  int64
+	// DecodeFields lets document backends avoid recursively decoding fields that
+	// the handler's inclusion projection, filter and sort cannot observe. Other
+	// backends may ignore it and return complete documents.
+	DecodeFields []string
 
 	OnlyRecordIDs bool
 	Comment       string

internal/backends/sqlite/collection.go

@@ -94,7 +94,7 @@ func (c *collection) Query(ctx context.Context, params *backends.QueryParams) (*
 		return nil, lazyerrors.Error(err)
 	}
 
-	iter := newQueryIterator(ctx, rows, params.OnlyRecordIDs)
+	iter := newQueryIterator(ctx, rows, params.OnlyRecordIDs, params.DecodeFields)
 	if os.Getenv("DEBUGSPEED") == "true" {
 		iter = newSpeedQueryIterator(iter, &querySpeed{
 			logger:        c.r.Logger(),

internal/backends/sqlite/query_iterator.go

@@ -41,6 +41,7 @@ type queryIterator struct {
 	token         *resource.Token
 	m             sync.Mutex
 	onlyRecordIDs bool
+	decodeFields  []string
 	speed         *querySpeed
 }
 
@@ -80,13 +81,16 @@ func newSpeedQueryIterator(iter *queryIterator, speed *querySpeed) *queryIterato
 //
 // Nil rows are possible and return already done iterator.
 // It still should be Close'd.
-func newQueryIterator(ctx context.Context, rows *fsql.Rows, onlyRecordIDs bool) *queryIterator {
+func newQueryIterator(ctx context.Context, rows *fsql.Rows, onlyRecordIDs bool, decodeFields ...[]string) *queryIterator {
 	iter := &queryIterator{
 		ctx:           ctx,
 		rows:          rows,
 		onlyRecordIDs: onlyRecordIDs,
 		token:         resource.NewToken(),
 	}
+	if len(decodeFields) != 0 {
+		iter.decodeFields = decodeFields[0]
+	}
 	resource.Track(iter, iter.token)
 
 	return iter
@@ -154,7 +158,12 @@ func (iter *queryIterator) Next() (struct{}, *types.Document, error) {
 
 	if !iter.onlyRecordIDs {
 		decodeStarted := time.Now()
-		if doc, err = sjson.Unmarshal(b); err != nil {
+		if len(iter.decodeFields) == 0 {
+			doc, err = sjson.Unmarshal(b)
+		} else {
+			doc, err = sjson.UnmarshalFields(b, iter.decodeFields)
+		}
+		if err != nil {
 			iter.close()
 			return unused, nil, lazyerrors.Error(err)
 		}

internal/handler/msg_find.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"log/slog"
 	"math"
+	"slices"
 	"strings"
 	"time"
 
@@ -244,6 +245,18 @@ func (h *Handler) makeFindQueryParams(ctx context.Context, params *common.FindPa
 		Comment: params.Comment,
 	}
 
+	if _, inclusion, projectionErr := common.ValidateProjection(params.Projection); projectionErr == nil && inclusion {
+		fields := make(map[string]struct{})
+		collectDecodeFields(params.Projection, fields)
+		collectDecodeFields(params.Filter, fields)
+		collectDecodeFields(params.Sort, fields)
+		qp.DecodeFields = make([]string, 0, len(fields))
+		for field := range fields {
+			qp.DecodeFields = append(qp.DecodeFields, field)
+		}
+		slices.Sort(qp.DecodeFields)
+	}
+
 	var err error
 	if params.Filter != nil {
 		if qp.Comment, err = common.GetOptionalParam(params.Filter, "$comment", qp.Comment); err != nil {
@@ -316,6 +329,34 @@ func (h *Handler) makeFindQueryParams(ctx context.Context, params *common.FindPa
 	return qp, nil
 }
 
+// collectDecodeFields records top-level document fields observable by a query.
+// Operator documents are traversed; dotted paths require decoding their root.
+func collectDecodeFields(doc *types.Document, fields map[string]struct{}) {
+	if doc == nil {
+		return
+	}
+	for _, key := range doc.Keys() {
+		value := must.NotFail(doc.Get(key))
+		if strings.HasPrefix(key, "$") {
+			switch value := value.(type) {
+			case *types.Document:
+				collectDecodeFields(value, fields)
+			case *types.Array:
+				for i := 0; i < value.Len(); i++ {
+					if branch, ok := must.NotFail(value.Get(i)).(*types.Document); ok {
+						collectDecodeFields(branch, fields)
+					}
+				}
+			}
+			continue
+		}
+		root, _, _ := strings.Cut(strings.TrimSuffix(key, ".$"), ".")
+		if root != "" {
+			fields[root] = struct{}{}
+		}
+	}
+}
+
 // makeFindIter creates an iterator chain for the find command.
 //
 // Iter is passed from the backend's query.

internal/handler/msg_find_decode_test.go

@@ -0,0 +1,43 @@
+// Copyright 2021 FerretDB Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package handler
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/FerretDB/FerretDB/internal/types"
+	"github.com/FerretDB/FerretDB/internal/util/must"
+)
+
+func TestCollectDecodeFields(t *testing.T) {
+	t.Parallel()
+
+	fields := make(map[string]struct{})
+	collectDecodeFields(must.NotFail(types.NewDocument(
+		"_id", int64(1),
+		"profile.name", int64(1),
+		"services.resume.tokens.$", int64(1),
+		"$and", must.NotFail(types.NewArray(
+			must.NotFail(types.NewDocument("archived", false)),
+			must.NotFail(types.NewDocument("members.userId", "u1")),
+		)),
+	)), fields)
+
+	assert.Equal(t, map[string]struct{}{
+		"_id": {}, "profile": {}, "services": {}, "archived": {}, "members": {},
+	}, fields)
+}

internal/handler/sjson/sjson.go

@@ -237,6 +237,58 @@ func Unmarshal(data []byte) (*types.Document, error) {
 	return d, nil
 }
 
+// UnmarshalFields decodes only selected top-level fields. It still parses the
+// complete outer JSON object, but avoids recursively decoding unobservable
+// values and their often much larger schemas (for example member arrays when a
+// client requested only _id). Missing selected fields remain missing.
+func UnmarshalFields(data []byte, fields []string) (*types.Document, error) {
+	var values map[string]json.RawMessage
+	if err := json.Unmarshal(data, &values); err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	jsch, ok := values["$s"]
+	if !ok {
+		return nil, lazyerrors.Errorf("schema is not set")
+	}
+	var sch struct {
+		Properties map[string]json.RawMessage `json:"p"`
+		Keys       []string                   `json:"$k"`
+	}
+	if err := json.Unmarshal(jsch, &sch); err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	wanted := make(map[string]struct{}, len(fields))
+	for _, field := range fields {
+		wanted[field] = struct{}{}
+	}
+	d := must.NotFail(types.NewDocument())
+	for _, key := range sch.Keys {
+		if _, ok := wanted[key]; !ok {
+			continue
+		}
+		value, exists := values[key]
+		if !exists {
+			return nil, lazyerrors.Errorf("sjson.UnmarshalFields: missing key %q", key)
+		}
+		rawElem, exists := sch.Properties[key]
+		if !exists {
+			return nil, lazyerrors.Errorf("sjson.UnmarshalFields: missing schema for key %q", key)
+		}
+		var fieldSchema elem
+		if err := json.Unmarshal(rawElem, &fieldSchema); err != nil {
+			return nil, lazyerrors.Error(err)
+		}
+		decoded, err := unmarshalSingleValue(value, &fieldSchema)
+		if err != nil {
+			return nil, lazyerrors.Error(err)
+		}
+		d.Set(key, decoded)
+	}
+	return d, nil
+}
+
 // unmarshalSingleValue decodes the given sjson-encoded data element by the given schema.
 func unmarshalSingleValue(data json.RawMessage, sch *elem) (any, error) {
 	if bytes.Equal(data, []byte("null")) {

internal/handler/sjson/sjson_marshal_unmarshal_test.go

@@ -66,6 +66,32 @@ func TestMarshalUnmarshal(t *testing.T) {
 	}
 }
 
+func TestUnmarshalFields(t *testing.T) {
+	t.Parallel()
+
+	nested := must.NotFail(types.NewDocument(
+		"members", must.NotFail(types.NewArray(
+			must.NotFail(types.NewDocument("userId", "u1", "active", true)),
+		)),
+	))
+	doc := must.NotFail(types.NewDocument(
+		"title", "large board",
+		"_id", types.ObjectID{1, 2, 3},
+		"archived", false,
+		"nested", nested,
+	))
+	b, err := Marshal(doc)
+	require.NoError(t, err)
+
+	actual, err := UnmarshalFields(b, []string{"archived", "_id", "missing"})
+	require.NoError(t, err)
+	assert.Equal(t, must.NotFail(types.NewDocument(
+		"_id", types.ObjectID{1, 2, 3},
+		"archived", false,
+	)), actual, "selected fields retain stored order and missing fields remain missing")
+	assert.False(t, actual.Has("nested"), "unobservable nested values must not be recursively decoded")
+}
+
 // TestUnmarshalInvalid checks that in case of invalid data, we return errors and not just ignore issues.
 func TestUnmarshalInvalid(t *testing.T) {
 	t.Parallel()