CSHARP-5321: Optimize client-side projections to only fetch the parts of the document that are needed for the projection.

src/MongoDB.Driver/ClientSideProjectionSnippetsDeserializer.cs

@@ -0,0 +1,81 @@
+/* Copyright 2010-present MongoDB 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.
+*/
+
+using System;
+using MongoDB.Bson;
+using MongoDB.Bson.IO;
+using MongoDB.Bson.Serialization;
+using MongoDB.Bson.Serialization.Serializers;
+
+namespace MongoDB.Driver
+{
+    internal static class ClientSideProjectionSnippetsDeserializer
+    {
+        public static IBsonSerializer Create(
+            Type projectionType,
+            IBsonSerializer[] snippetDeserializers,
+            Delegate projector)
+        {
+            var deserializerType = typeof(ClientSideProjectionSnippetsDeserializer<>).MakeGenericType(projectionType);
+            return (IBsonSerializer)Activator.CreateInstance(deserializerType, [snippetDeserializers, projector]);
+        }
+    }
+
+    internal sealed class ClientSideProjectionSnippetsDeserializer<TProjection> : SerializerBase<TProjection>, IClientSideProjectionDeserializer
+    {
+        private readonly IBsonSerializer[] _snippetDeserializers;
+        private readonly Func<object[], TProjection> _projector;
+
+        public ClientSideProjectionSnippetsDeserializer(IBsonSerializer[] snippetDeserializers, Func<object[], TProjection> projector)
+        {
+            _snippetDeserializers = snippetDeserializers;
+            _projector = projector;
+        }
+
+        public override TProjection Deserialize(BsonDeserializationContext context, BsonDeserializationArgs args)
+        {
+            var snippets = DeserializeSnippets(context);
+            return _projector(snippets);
+        }
+
+        private object[] DeserializeSnippets(BsonDeserializationContext context)
+        {
+            var reader = context.Reader;
+
+            reader.ReadStartDocument();
+            reader.ReadName("_snippets");
+            reader.ReadStartArray();
+            var snippets = new object[_snippetDeserializers.Length];
+            var i = 0;
+            while (reader.ReadBsonType() != BsonType.EndOfDocument)
+            {
+                if (i >= _snippetDeserializers.Length)
+                {
+                    throw new BsonSerializationException($"Expected {_snippetDeserializers.Length} snippets but found more than that.");
+                }
+                snippets[i] = _snippetDeserializers[i].Deserialize(context);
+                i++;
+            }
+            if (i != _snippetDeserializers.Length)
+            {
+                throw new BsonSerializationException($"Expected {_snippetDeserializers.Length} snippets but found {i}.");
+            }
+            reader.ReadEndArray();
+            reader.ReadEndDocument();
+
+            return snippets;
+        }
+    }
+}

src/MongoDB.Driver/Linq/Linq3Implementation/Misc/ExpressionIsReferencedVisitor.cs

@@ -20,6 +20,15 @@ namespace MongoDB.Driver.Linq.Linq3Implementation.Misc
 {
     internal class ExpressionIsReferencedVisitor : ExpressionVisitor
     {
+        #region static
+        public static bool IsReferenced(Expression node, Expression expression)
+        {
+            var visitor = new ExpressionIsReferencedVisitor(expression);
+            visitor.Visit(node);
+            return visitor.ExpressionIsReferenced;
+        }
+        #endregion
+
         private readonly Expression _expression;
         private bool _expressionIsReferenced;
 

src/MongoDB.Driver/Linq/Linq3Implementation/Misc/LambdaExpressionExtensions.cs

@@ -26,9 +26,7 @@ internal static class LambdaExpressionExtensions
     {
         public static bool LambdaBodyReferencesParameter(this LambdaExpression lambda, ParameterExpression parameter)
         {
-            var visitor = new ExpressionIsReferencedVisitor(parameter);
-            visitor.Visit(lambda.Body);
-            return visitor.ExpressionIsReferenced;
+            return ExpressionIsReferencedVisitor.IsReferenced(lambda.Body, parameter);
         }
 
         public static string TranslateToDottedFieldName(this LambdaExpression fieldSelectorLambda, TranslationContext context, IBsonSerializer parameterSerializer)

src/MongoDB.Driver/Linq/Linq3Implementation/Translators/ExpressionToAggregationExpressionTranslators/ClientSideProjectionRewriter.cs

@@ -0,0 +1,154 @@
+/* Copyright 2010-present MongoDB 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.
+ */
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Linq.Expressions;
+using System.Reflection;
+using MongoDB.Bson.Serialization;
+using MongoDB.Driver.Linq.Linq3Implementation.Misc;
+using MongoDB.Driver.Linq.Linq3Implementation.Reflection;
+using ExpressionVisitor = System.Linq.Expressions.ExpressionVisitor;
+
+namespace MongoDB.Driver.Linq.Linq3Implementation.Translators.ExpressionToAggregationExpressionTranslators
+{
+    internal class ClientSideProjectionRewriter: ExpressionVisitor
+    {
+        #region static
+        private readonly static MethodInfo[] __orderByMethods =
+        [
+            EnumerableMethod.OrderBy,
+            EnumerableMethod.OrderByDescending,
+            QueryableMethod.OrderBy,
+            QueryableMethod.OrderByDescending
+        ];
+
+        private readonly static MethodInfo[] __thenByMethods =
+        [
+            EnumerableMethod.ThenBy,
+            EnumerableMethod.ThenByDescending,
+            QueryableMethod.ThenBy,
+            QueryableMethod.ThenByDescending
+        ];
+
+        public static (TranslatedExpression[], LambdaExpression) RewriteProjection(TranslationContext context, LambdaExpression projectionLambda, IBsonSerializer sourceSerializer)
+        {
+            var rootParameter = projectionLambda.Parameters.Single();
+            var rootSymbol = context.CreateRootSymbol(rootParameter, sourceSerializer);
+            context = context.WithSymbol(rootSymbol);
+
+            var snippetsParameter = Expression.Parameter(typeof(object[]), "snippets");
+            var projectionRewriter = new ClientSideProjectionRewriter(context, snippetsParameter);
+            var rewrittenBody = projectionRewriter.Visit(projectionLambda.Body);
+            var rewrittenLambda = Expression.Lambda(rewrittenBody, snippetsParameter);
+            var snippetsArray = projectionRewriter.Snippets.ToArray();
+
+            return (snippetsArray, rewrittenLambda);
+        }
+        #endregion
+
+        private readonly TranslationContext _context;
+        private readonly List<TranslatedExpression> _snippets = new();
+        private readonly ParameterExpression _snippetsParameter;
+
+        private ClientSideProjectionRewriter(TranslationContext context, ParameterExpression snippetsParameter)
+        {
+            _context = context;
+            _snippetsParameter = snippetsParameter;
+        }
+
+        private List<TranslatedExpression> Snippets => _snippets;
+
+        public override Expression Visit(Expression node)
+        {
+            if (node == null)
+            {
+                return null;
+            }
+
+            if (node.NodeType == ExpressionType.Constant)
+            {
+                return node; // don't make snippets for constants
+            }
+
+            TranslatedExpression snippet;
+            try
+            {
+                snippet = ExpressionToAggregationExpressionTranslator.Translate(_context, node);
+            }
+            catch
+            {
+                return base.Visit(node); // try to find smaller snippets below this node
+            }
+
+            var snippetIndex = _snippets.Count;
+            _snippets.Add(snippet);
+
+            var snippetReference = // (T)snippets[i]
+                Expression.Convert(
+                    Expression.ArrayIndex(_snippetsParameter, Expression.Constant(snippetIndex)),
+                    snippet.Expression.Type);
+
+            return snippetReference;
+        }
+
+        protected override Expression VisitMethodCall(MethodCallExpression node)
+        {
+            // don't split OrderBy/ThenBy across the client/server boundary
+            if (node.Method.IsOneOf(__thenByMethods))
+            {
+                return VisitThenBy(node);
+            }
+
+            return base.VisitMethodCall(node);
+        }
+
+        private Expression VisitThenBy(MethodCallExpression node)
+        {
+            var arguments = node.Arguments;
+            var sourceExpression = arguments[0];
+            var keySelectorExpression = arguments[1];
+
+            if (sourceExpression is MethodCallExpression sourceMethodCallExpression)
+            {
+                var sourceMethod = sourceMethodCallExpression.Method;
+
+                if (sourceMethod.IsOneOf(__thenByMethods))
+                {
+                    var rewrittenSourceExpression = VisitThenBy(sourceMethodCallExpression);
+                    return node.Update(node.Object, [rewrittenSourceExpression, keySelectorExpression]);
+                }
+
+                if (sourceMethod.IsOneOf(__orderByMethods))
+                {
+                    var rewrittenSourceExpression = VisitOrderBy(sourceMethodCallExpression);
+                    return node.Update(node.Object, [rewrittenSourceExpression, keySelectorExpression]);
+                }
+            }
+
+            throw new ArgumentException("ThenBy or ThenByDescending not preceded by OrderBy or OrderByDescending.", nameof(node));
+        }
+
+        private Expression VisitOrderBy(MethodCallExpression node)
+        {
+            var arguments = node.Arguments;
+            var sourceExpression = arguments[0];
+            var keySelectorExpression = arguments[1];
+            var rewrittenSourceExpression = Visit(sourceExpression);
+            return node.Update(node.Object, [rewrittenSourceExpression, keySelectorExpression]);
+        }
+    }
+}

src/MongoDB.Driver/Linq/Linq3Implementation/Translators/ExpressionToAggregationExpressionTranslators/ClientSideProjectionTranslator.cs

@@ -0,0 +1,75 @@
+/* Copyright 2010-present MongoDB 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.
+ */
+
+using System.Linq;
+using System.Linq.Expressions;
+using MongoDB.Bson.Serialization;
+using MongoDB.Driver.Linq.Linq3Implementation.Ast.Expressions;
+using MongoDB.Driver.Linq.Linq3Implementation.Ast.Stages;
+using MongoDB.Driver.Linq.Linq3Implementation.Misc;
+
+namespace MongoDB.Driver.Linq.Linq3Implementation.Translators.ExpressionToAggregationExpressionTranslators
+{
+    internal static class ClientSideProjectionTranslator
+    {
+        public static (AstProjectStage, IBsonSerializer) CreateProjectSnippetsStage(
+            TranslationContext context,
+            LambdaExpression projectionLambda,
+            IBsonSerializer sourceSerializer)
+        {
+            var (snippetsAst, snippetsProjectionDeserializer) = RewriteProjectionUsingSnippets(context, projectionLambda, sourceSerializer);
+            if (snippetsAst == null)
+            {
+                return (null, snippetsProjectionDeserializer);
+            }
+            else
+            {
+                var snippetsTranslation = new TranslatedExpression(projectionLambda, snippetsAst, snippetsProjectionDeserializer);
+                return ProjectionHelper.CreateProjectStage(snippetsTranslation);
+            }
+        }
+
+        private static (AstComputedDocumentExpression, IBsonSerializer) RewriteProjectionUsingSnippets(
+            TranslationContext context,
+            LambdaExpression projectionLambda,
+            IBsonSerializer sourceSerializer)
+        {
+            var (snippets, rewrittenProjectionLamdba) = ClientSideProjectionRewriter.RewriteProjection(context, projectionLambda, sourceSerializer);
+
+            if (snippets.Length == 0 || snippets.Any(IsRoot))
+            {
+                var clientSideProjectionDeserializer = ClientSideProjectionDeserializer.Create(sourceSerializer, projectionLambda);
+                return (null, clientSideProjectionDeserializer); // project directly off $$ROOT with no snippets
+            }
+            else
+            {
+                var snippetsComputedDocument = CreateSnippetsComputedDocument(snippets);
+                var snippetDeserializers = snippets.Select(s => s.Serializer).ToArray();
+                var rewrittenProjectionDelegate = rewrittenProjectionLamdba.Compile();
+                var clientSideProjectionSnippetsDeserializer = ClientSideProjectionSnippetsDeserializer.Create(projectionLambda.ReturnType, snippetDeserializers, rewrittenProjectionDelegate);
+                return (snippetsComputedDocument, clientSideProjectionSnippetsDeserializer);
+            }
+
+            static bool IsRoot(TranslatedExpression snippet) => snippet.Ast.IsRootVar();
+        }
+
+        private static AstComputedDocumentExpression CreateSnippetsComputedDocument(TranslatedExpression[] snippets)
+        {
+            var snippetsArray = AstExpression.ComputedArray(snippets.Select(s => s.Ast));
+            var snippetsField = AstExpression.ComputedField("_snippets", snippetsArray);
+            return (AstComputedDocumentExpression)AstExpression.ComputedDocument([snippetsField]);
+        }
+    }
+}

src/MongoDB.Driver/Linq/Linq3Implementation/Translators/ExpressionToPipelineTranslators/SelectMethodToPipelineTranslator.cs

@@ -15,6 +15,7 @@
 
 using System;
 using System.Linq.Expressions;
+using MongoDB.Bson.Serialization;
 using MongoDB.Driver.Linq.Linq3Implementation.Ast;
 using MongoDB.Driver.Linq.Linq3Implementation.Ast.Stages;
 using MongoDB.Driver.Linq.Linq3Implementation.Misc;
@@ -46,19 +47,21 @@ public static TranslatedPipeline Translate(TranslationContext context, MethodCal
                 ClientSideProjectionHelper.ThrowIfClientSideProjection(expression, pipeline, method);
 
                 var sourceSerializer = pipeline.OutputSerializer;
+                AstProjectStage projectStage;
+                IBsonSerializer projectionSerializer;
                 try
                 {
                     var selectorTranslation = ExpressionToAggregationExpressionTranslator.TranslateLambdaBody(context, selectorLambda, sourceSerializer, asRoot: true);
-                    var (projectStage, projectionSerializer) = ProjectionHelper.CreateProjectStage(selectorTranslation);
-                    pipeline = pipeline.AddStage(projectStage, projectionSerializer);
+                    (projectStage, projectionSerializer) = ProjectionHelper.CreateProjectStage(selectorTranslation);
                 }
                 catch (ExpressionNotSupportedException) when (context.TranslationOptions?.EnableClientSideProjections ?? false)
                 {
-                    var clientSideProjectionDeserializer = ClientSideProjectionDeserializer.Create(sourceSerializer, selectorLambda);
-                    pipeline = pipeline.WithNewOutputSerializer(clientSideProjectionDeserializer);
+                    (projectStage, projectionSerializer) = ClientSideProjectionTranslator.CreateProjectSnippetsStage(context, selectorLambda, sourceSerializer);
                 }
 
-                return pipeline;
+                return projectStage == null ?
+                    pipeline.WithNewOutputSerializer(projectionSerializer) : // project directly off $$ROOT with no $project stage
+                    pipeline.AddStage(projectStage, projectionSerializer);
             }
 
             throw new ExpressionNotSupportedException(expression);