Replace document slices with iterators

internal/handlers/common/distinct.go

@@ -15,6 +15,7 @@
 package common
 
 import (
+	"errors"
 	"fmt"
 	"strings"
 
@@ -23,6 +24,7 @@
 	"github.com/FerretDB/FerretDB/internal/handlers/commonerrors"
 	"github.com/FerretDB/FerretDB/internal/handlers/commonparams"
 	"github.com/FerretDB/FerretDB/internal/types"
+	"github.com/FerretDB/FerretDB/internal/util/iterator"
 	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
 )
 
@@ -84,10 +86,21 @@
 //
 // If the key is found in the document, and the value is an array, each element of the array is added to the result.
 // Otherwise, the value itself is added to the result.
-func FilterDistinctValues(docs []*types.Document, key string) (*types.Array, error) {
-	distinct := types.MakeArray(len(docs))
+func FilterDistinctValues(iter types.DocumentsIterator, key string) (*types.Array, error) {
+	distinct := types.MakeArray(0)
+
+	defer iter.Close()
+
+	for {
+		_, doc, err := iter.Next()
+		if errors.Is(err, iterator.ErrIteratorDone) {
+			break
+		}
+
+		if err != nil {
+			return nil, lazyerrors.Error(err)
+		}
 
-	for _, doc := range docs {
 		// docsAtSuffix contains all documents exist at the suffix key.
 		docsAtSuffix := []*types.Document{doc}
 		suffix := key

internal/handlers/pg/msg_count.go

@@ -22,6 +22,7 @@
 	"github.com/FerretDB/FerretDB/internal/handlers/common"
 	"github.com/FerretDB/FerretDB/internal/handlers/pg/pgdb"
 	"github.com/FerretDB/FerretDB/internal/types"
+	"github.com/FerretDB/FerretDB/internal/util/iterator"
 	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
 	"github.com/FerretDB/FerretDB/internal/util/must"
 	"github.com/FerretDB/FerretDB/internal/wire"
@@ -50,26 +51,40 @@
 		Collection: params.Collection,
 	}
 
-	qp.Filter = params.Filter
+	if !h.DisableFilterPushdown {
+		qp.Filter = params.Filter
+	}
 
 	var resDocs []*types.Document
 	err = dbPool.InTransaction(ctx, func(tx pgx.Tx) error {
-		resDocs, err = fetchAndFilterDocs(ctx, &fetchParams{tx, &qp, h.DisableFilterPushdown})
-		return err
+		var iter types.DocumentsIterator
+
+		iter, _, err = pgdb.QueryDocuments(ctx, tx, &qp)
+		if err != nil {
+			return lazyerrors.Error(err)
+		}
+
+		closer := iterator.NewMultiCloser(iter)
+		defer closer.Close()
+
+		iter = common.FilterIterator(iter, closer, qp.Filter)
+
+		iter = common.SkipIterator(iter, closer, params.Skip)
+
+		iter = common.LimitIterator(iter, closer, params.Limit)
+
+		resDocs, err = iterator.ConsumeValues(iterator.Interface[struct{}, *types.Document](iter))
+		if err != nil {
+			return lazyerrors.Error(err)
+		}
+
+		return nil
 	})
 
 	if err != nil {
 		return nil, err
 	}
 
-	if resDocs, err = common.SkipDocuments(resDocs, params.Skip); err != nil {
-		return nil, lazyerrors.Error(err)
-	}
-
-	if resDocs, err = common.LimitDocuments(resDocs, params.Limit); err != nil {
-		return nil, lazyerrors.Error(err)
-	}
-
 	var reply wire.OpMsg
 	must.NoError(reply.SetSections(wire.OpMsgSection{
 		Documents: []*types.Document{must.NotFail(types.NewDocument(

internal/handlers/pg/msg_distinct.go

@@ -22,6 +22,7 @@
 	"github.com/FerretDB/FerretDB/internal/handlers/common"
 	"github.com/FerretDB/FerretDB/internal/handlers/pg/pgdb"
 	"github.com/FerretDB/FerretDB/internal/types"
+	"github.com/FerretDB/FerretDB/internal/util/iterator"
 	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
 	"github.com/FerretDB/FerretDB/internal/util/must"
 	"github.com/FerretDB/FerretDB/internal/wire"
@@ -47,21 +48,35 @@
 	qp := pgdb.QueryParams{
 		DB:         dp.DB,
 		Collection: dp.Collection,
-		Filter:     dp.Filter,
 		Comment:    dp.Comment,
 	}
 
-	var resDocs []*types.Document
+	if !h.DisableFilterPushdown {
+		qp.Filter = dp.Filter
+	}
+
+	var distinct *types.Array
 	err = dbPool.InTransaction(ctx, func(tx pgx.Tx) error {
-		resDocs, err = fetchAndFilterDocs(ctx, &fetchParams{tx, &qp, h.DisableFilterPushdown})
-		return err
-	})
+		var iter types.DocumentsIterator
 
-	if err != nil {
-		return nil, err
-	}
+		iter, _, err = pgdb.QueryDocuments(ctx, tx, &qp)
+		if err != nil {
+			return lazyerrors.Error(err)
+		}
+
+		closer := iterator.NewMultiCloser(iter)
+		defer closer.Close()
+
+		iter = common.FilterIterator(iter, closer, qp.Filter)
+
+		distinct, err = common.FilterDistinctValues(iter, dp.Key)
+		if err != nil {
+			return lazyerrors.Error(err)
+		}
+
+		return nil
+	})
 
-	distinct, err := common.FilterDistinctValues(resDocs, dp.Key)
 	if err != nil {
 		return nil, err
 	}
@@ -74,9 +89,5 @@
 		))},
 	}))
 
-	if err != nil {
-		return nil, lazyerrors.Error(err)
-	}
-
 	return &reply, nil
 }

internal/handlers/tigris/msg_distinct.go

@@ -20,6 +20,7 @@
 	"github.com/FerretDB/FerretDB/internal/handlers/common"
 	"github.com/FerretDB/FerretDB/internal/handlers/tigris/tigrisdb"
 	"github.com/FerretDB/FerretDB/internal/types"
+	"github.com/FerretDB/FerretDB/internal/util/iterator"
 	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
 	"github.com/FerretDB/FerretDB/internal/util/must"
 	"github.com/FerretDB/FerretDB/internal/wire"
@@ -42,18 +43,28 @@
 		return nil, err
 	}
 
+	filter := dp.Filter
+
 	qp := tigrisdb.QueryParams{
 		DB:         dp.DB,
 		Collection: dp.Collection,
-		Filter:     dp.Filter,
 	}
 
-	resDocs, err := fetchAndFilterDocs(ctx, &fetchParams{dbPool, &qp, h.DisableFilterPushdown})
+	if !h.DisableFilterPushdown {
+		qp.Filter = filter
+	}
+
+	iter, err := dbPool.QueryDocuments(ctx, &qp)
 	if err != nil {
-		return nil, err
+		return nil, lazyerrors.Error(err)
 	}
 
-	distinct, err := common.FilterDistinctValues(resDocs, dp.Key)
+	closer := iterator.NewMultiCloser(iter)
+	defer closer.Close()
+
+	iter = common.FilterIterator(iter, closer, filter)
+
+	distinct, err := common.FilterDistinctValues(iter, dp.Key)
 	if err != nil {
 		return nil, err
 	}