Fix to #17066 - Query: Translate ElementAt(OrDefault)

Enable member pushdown for ElementAt(OrDefault) and then handle the method in translation in a similar way to First(OrDefault) with an optional Skip(x) element.

Fixes #17066

src/EFCore.Relational/Query/RelationalQueryableMethodTranslatingExpressionVisitor.cs

@@ -395,7 +395,24 @@ private static ShapedQueryExpression CreateShapedQueryExpression(IEntityType ent
         ShapedQueryExpression source,
         Expression index,
         bool returnDefault)
-        => null;
+    {
+        var selectExpression = (SelectExpression)source.QueryExpression;
+        var translation = TranslateExpression(index);
+        if (translation == null)
+        {
+            return null;
+        }
+
+        if (selectExpression.Orderings.Count == 0)
+        {
+            _queryCompilationContext.Logger.RowLimitingOperationWithoutOrderByWarning();
+        }
+
+        selectExpression.ApplyOffset(translation);
+        selectExpression.ApplyLimit(TranslateExpression(Expression.Constant(1))!);
+
+        return source;
+    }
 
     /// <inheritdoc />
     protected override ShapedQueryExpression? TranslateExcept(ShapedQueryExpression source1, ShapedQueryExpression source2)

src/EFCore.Relational/Query/RelationalSqlTranslatingExpressionVisitor.cs

@@ -35,9 +35,9 @@ public class RelationalSqlTranslatingExpressionVisitor : ExpressionVisitor
         QueryableMethods.LastWithPredicate,
         QueryableMethods.LastWithoutPredicate,
         QueryableMethods.LastOrDefaultWithPredicate,
-        QueryableMethods.LastOrDefaultWithoutPredicate
-        //QueryableMethodProvider.ElementAtMethodInfo,
-        //QueryableMethodProvider.ElementAtOrDefaultMethodInfo
+        QueryableMethods.LastOrDefaultWithoutPredicate,
+        QueryableMethods.ElementAt,
+        QueryableMethods.ElementAtOrDefault
     };
 
     private static readonly List<MethodInfo> PredicateAggregateMethodInfos = new()
@@ -346,13 +346,29 @@ protected override Expression VisitBinary(BinaryExpression binaryExpression)
                 && SingleResultMethodInfos.Contains(nonNullMethodCallExpression.Method.GetGenericMethodDefinition()))
             {
                 var source = nonNullMethodCallExpression.Arguments[0];
-                if (nonNullMethodCallExpression.Arguments.Count == 2)
+                var genericMethod = nonNullMethodCallExpression.Method.GetGenericMethodDefinition();
+                if (genericMethod == QueryableMethods.FirstWithPredicate
+                    || genericMethod == QueryableMethods.FirstOrDefaultWithPredicate
+                    || genericMethod == QueryableMethods.SingleWithPredicate
+                    || genericMethod == QueryableMethods.SingleOrDefaultWithPredicate
+                    || genericMethod == QueryableMethods.LastWithPredicate
+                    || genericMethod == QueryableMethods.LastOrDefaultWithPredicate)
                 {
                     source = Expression.Call(
                         QueryableMethods.Where.MakeGenericMethod(source.Type.GetSequenceType()),
                         source,
                         nonNullMethodCallExpression.Arguments[1]);
                 }
+                else if ((genericMethod == QueryableMethods.ElementAt || genericMethod == QueryableMethods.ElementAtOrDefault)
+                    && (nonNullMethodCallExpression.Arguments[1] is not ConstantExpression constantIndex
+                        || constantIndex.Value is not int constantInt
+                        || constantInt != 0))
+                {
+                    source = Expression.Call(
+                        QueryableMethods.Skip.MakeGenericMethod(source.Type.GetSequenceType()),
+                        source,
+                        nonNullMethodCallExpression.Arguments[1]);
+                }
 
                 var translatedSubquery = _queryableMethodTranslatingExpressionVisitor.TranslateSubquery(source);
                 if (translatedSubquery != null)

src/EFCore/Extensions/EntityFrameworkQueryableExtensions.cs

@@ -275,6 +275,87 @@ public static Task<long> LongCountAsync<TSource>(
 
     #endregion
 
+    #region ElementAt/ElementAtOrDefault
+
+    /// <summary>
+    ///     Asynchronously returns the element at a specified index in a sequence.
+    /// </summary>
+    /// <remarks>
+    ///     <para>
+    ///         Multiple active operations on the same context instance are not supported. Use <see langword="await" /> to ensure
+    ///         that any asynchronous operations have completed before calling another method on this context.
+    ///         See <see href="https://aka.ms/efcore-docs-threading">Avoiding DbContext threading issues</see> for more information and examples.
+    ///     </para>
+    ///     <para>
+    ///         See <see href="https://aka.ms/efcore-docs-async-linq">Querying data with EF Core</see> for more information and examples.
+    ///     </para>
+    /// </remarks>
+    /// <typeparam name="TSource">The type of the elements of <paramref name="source" />.</typeparam>
+    /// <param name="source">An <see cref="IQueryable{T}" /> to return the element from.</param>
+    /// <param name="index">The zero-based index of the element to retrieve.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken" /> to observe while waiting for the task to complete.</param>
+    /// <returns>
+    ///     A task that represents the asynchronous operation.
+    ///     The task result contains the element at a specified index in a <paramref name="source" /> sequence.
+    /// </returns>
+    /// <exception cref="ArgumentNullException">
+    ///     <paramref name="source" /> is <see langword="null" />.
+    /// </exception>
+    /// <exception cref="ArgumentOutOfRangeException">
+    ///     <para>
+    ///         <paramref name="index" /> is less than zero.
+    ///     </para>
+    /// </exception>
+    /// <exception cref="OperationCanceledException">If the <see cref="CancellationToken" /> is canceled.</exception>
+    public static Task<TSource> ElementAtAsync<TSource>(
+        this IQueryable<TSource> source,
+        int index,
+        CancellationToken cancellationToken = default)
+    {
+        Check.NotNull(index, nameof(index));
+
+        return ExecuteAsync<TSource, Task<TSource>>(
+            QueryableMethods.ElementAt, source, Expression.Constant(index), cancellationToken);
+    }
+
+    /// <summary>
+    ///     Asynchronously returns the element at a specified index in a sequence, or a default value if the index is out of range.
+    /// </summary>
+    /// <remarks>
+    ///     <para>
+    ///         Multiple active operations on the same context instance are not supported. Use <see langword="await" /> to ensure
+    ///         that any asynchronous operations have completed before calling another method on this context.
+    ///         See <see href="https://aka.ms/efcore-docs-threading">Avoiding DbContext threading issues</see> for more information and examples.
+    ///     </para>
+    ///     <para>
+    ///         See <see href="https://aka.ms/efcore-docs-async-linq">Querying data with EF Core</see> for more information and examples.
+    ///     </para>
+    /// </remarks>
+    /// <typeparam name="TSource">The type of the elements of <paramref name="source" />.</typeparam>
+    /// <param name="source">An <see cref="IQueryable{T}" /> to return the element from.</param>
+    /// <param name="index">The zero-based index of the element to retrieve.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken" /> to observe while waiting for the task to complete.</param>
+    /// <returns>
+    ///     A task that represents the asynchronous operation.
+    ///     The task result contains the element at a specified index in a <paramref name="source" /> sequence.
+    /// </returns>
+    /// <exception cref="ArgumentNullException">
+    ///     <paramref name="source" /> is <see langword="null" />.
+    /// </exception>
+    /// <exception cref="OperationCanceledException">If the <see cref="CancellationToken" /> is canceled.</exception>
+    public static Task<TSource> ElementAtOrDefaultAsync<TSource>(
+        this IQueryable<TSource> source,
+        int index,
+        CancellationToken cancellationToken = default)
+    {
+        Check.NotNull(index, nameof(index));
+
+        return ExecuteAsync<TSource, Task<TSource>>(
+            QueryableMethods.ElementAtOrDefault, source, Expression.Constant(index), cancellationToken);
+    }
+
+    #endregion
+
     #region First/FirstOrDefault
 
     /// <summary>

src/EFCore/Query/Internal/NavigationExpandingExpressionVisitor.ExpressionVisitors.cs

@@ -1000,10 +1000,13 @@ private sealed class ReducingExpressionVisitor : ExpressionVisitor
 
                     if (navigationExpansionExpression.CardinalityReducingGenericMethodInfo != null)
                     {
+                        var arguments = new List<Expression> { result };
+                        arguments.AddRange(navigationExpansionExpression.CardinalityReducingMethodArguments.Select(x => Visit(x)));
+
                         result = Expression.Call(
                             navigationExpansionExpression.CardinalityReducingGenericMethodInfo.MakeGenericMethod(
                                 result.Type.GetSequenceType()),
-                            result);
+                            arguments.ToArray());
                     }
 
                     return result;

src/EFCore/Query/Internal/NavigationExpandingExpressionVisitor.Expressions.cs

@@ -246,6 +246,7 @@ private set
 
         public Expression PendingSelector { get; private set; }
         public MethodInfo? CardinalityReducingGenericMethodInfo { get; private set; }
+        public List<Expression> CardinalityReducingMethodArguments { get; private set; } = new();
 
         public Type SourceElementType
             => CurrentParameter.Type;
@@ -274,8 +275,11 @@ public void AppendPendingOrdering(MethodInfo orderingMethod, Expression keySelec
         public void ClearPendingOrderings()
             => _pendingOrderings.Clear();
 
-        public void ConvertToSingleResult(MethodInfo genericMethod)
-            => CardinalityReducingGenericMethodInfo = genericMethod;
+        public void ConvertToSingleResult(MethodInfo genericMethod, params Expression[] arguments)
+        {
+            CardinalityReducingGenericMethodInfo = genericMethod;
+            CardinalityReducingMethodArguments.AddRange(arguments);
+        }
 
         public override ExpressionType NodeType
             => ExpressionType.Extension;

src/EFCore/Query/Internal/NavigationExpandingExpressionVisitor.cs

@@ -433,6 +433,16 @@ when QueryableMethods.IsSumWithSelector(method):
                             methodCallExpression.Arguments[1].UnwrapLambdaFromQuote(),
                             methodCallExpression.Type);
 
+                    case nameof(Queryable.ElementAt)
+                        when genericMethod == QueryableMethods.ElementAt:
+                    case nameof(Queryable.ElementAtOrDefault)
+                        when genericMethod == QueryableMethods.ElementAtOrDefault:
+                        return ProcessElementAt(
+                            source,
+                            genericMethod,
+                            methodCallExpression.Arguments[1],
+                            methodCallExpression.Type);
+
                     case nameof(Queryable.Join)
                         when genericMethod == QueryableMethods.Join:
                     {
@@ -982,6 +992,22 @@ private NavigationExpansionExpression ProcessFirstSingleLastOrDefault(
         return source;
     }
 
+    private NavigationExpansionExpression ProcessElementAt(
+        NavigationExpansionExpression source,
+        MethodInfo genericMethod,
+        Expression index,
+        Type returnType)
+    {
+        if (source.PendingSelector.Type != returnType)
+        {
+            source.ApplySelector(Expression.Convert(source.PendingSelector, returnType));
+        }
+
+        source.ConvertToSingleResult(genericMethod, index);
+
+        return source;
+    }
+
     // This returns Expression since it can also return a deferred GroupBy operation
     private Expression ProcessGroupBy(
         NavigationExpansionExpression source,

src/EFCore/Query/Internal/SubqueryMemberPushdownExpressionVisitor.cs

@@ -24,9 +24,9 @@ public class SubqueryMemberPushdownExpressionVisitor : ExpressionVisitor
         QueryableMethods.LastWithPredicate,
         QueryableMethods.LastWithoutPredicate,
         QueryableMethods.LastOrDefaultWithPredicate,
-        QueryableMethods.LastOrDefaultWithoutPredicate
-        //QueryableMethodProvider.ElementAtMethodInfo,
-        //QueryableMethodProvider.ElementAtOrDefaultMethodInfo
+        QueryableMethods.LastOrDefaultWithoutPredicate,
+        QueryableMethods.ElementAt,
+        QueryableMethods.ElementAtOrDefault
     };
 
     private static readonly IDictionary<MethodInfo, MethodInfo> PredicateLessMethodInfo = new Dictionary<MethodInfo, MethodInfo>
@@ -163,7 +163,13 @@ private Expression PushdownMember(
         var source = methodCallExpression.Arguments[0];
         var queryableType = source.Type.GetSequenceType();
         var genericMethod = methodCallExpression.Method.GetGenericMethodDefinition();
-        if (methodCallExpression.Arguments.Count == 2)
+
+        if (genericMethod == QueryableMethods.FirstWithPredicate
+            || genericMethod == QueryableMethods.FirstOrDefaultWithPredicate
+            || genericMethod == QueryableMethods.SingleWithPredicate
+            || genericMethod == QueryableMethods.SingleOrDefaultWithPredicate
+            || genericMethod == QueryableMethods.LastWithPredicate
+            || genericMethod == QueryableMethods.LastOrDefaultWithPredicate)
         {
             // Move predicate to Where so that we can change shape before operator
             source = Expression.Call(
@@ -203,7 +209,16 @@ private Expression PushdownMember(
                 Expression.Quote(Expression.Lambda(memberAccessExpression, parameter)));
         }
 
-        source = Expression.Call(genericMethod.MakeGenericMethod(source.Type.GetSequenceType()), source);
+        if (genericMethod == QueryableMethods.ElementAt
+            || genericMethod == QueryableMethods.ElementAtOrDefault)
+        {
+            var index = Visit(methodCallExpression.Arguments[1]);
+            source = Expression.Call(genericMethod.MakeGenericMethod(source.Type.GetSequenceType()), source, index);
+        }
+        else
+        {
+            source = Expression.Call(genericMethod.MakeGenericMethod(source.Type.GetSequenceType()), source);
+        }
 
         return source.Type != returnType
             ? Expression.Convert(source, returnType)

test/EFCore.Cosmos.FunctionalTests/Query/NorthwindMiscellaneousQueryCosmosTest.cs

@@ -4584,6 +4584,27 @@ public override async Task Select_subquery_recursive_trivial_returning_queryable
         AssertSql();
     }
 
+    public override async Task Collection_navigation_equal_to_null_for_subquery_using_ElementAtOrDefault_constant_zero(bool async)
+    {
+        await AssertTranslationFailed(() => base.Collection_navigation_equal_to_null_for_subquery_using_ElementAtOrDefault_constant_zero(async));
+
+        AssertSql();
+    }
+
+    public override async Task Collection_navigation_equal_to_null_for_subquery_using_ElementAtOrDefault_constant_one(bool async)
+    {
+        await AssertTranslationFailed(() => base.Collection_navigation_equal_to_null_for_subquery_using_ElementAtOrDefault_constant_one(async));
+
+        AssertSql();
+    }
+
+    public override async Task Collection_navigation_equal_to_null_for_subquery_using_ElementAtOrDefault_parameter(bool async)
+    {
+        await AssertTranslationFailed(() => base.Collection_navigation_equal_to_null_for_subquery_using_ElementAtOrDefault_parameter(async));
+
+        AssertSql();
+    }
+
     private void AssertSql(params string[] expected)
         => Fixture.TestSqlLoggerFactory.AssertBaseline(expected);
 

test/EFCore.Cosmos.FunctionalTests/Query/NorthwindWhereQueryCosmosTest.cs

@@ -2413,6 +2413,22 @@ public override async Task First_over_custom_projection_compared_to_not_null(boo
         AssertSql();
     }
 
+    public override async Task ElementAt_over_custom_projection_compared_to_not_null(bool async)
+    {
+        // Cosmos client evaluation. Issue #17246.
+        await AssertTranslationFailed(() => base.ElementAt_over_custom_projection_compared_to_not_null(async));
+
+        AssertSql();
+    }
+
+    public override async Task ElementAtOrDefault_over_custom_projection_compared_to_null(bool async)
+    {
+        // Cosmos client evaluation. Issue #17246.
+        await AssertTranslationFailed(() => base.ElementAtOrDefault_over_custom_projection_compared_to_null(async));
+
+        AssertSql();
+    }
+
     public override async Task Single_over_custom_projection_compared_to_null(bool async)
     {
         // Cosmos client evaluation. Issue #17246.

test/EFCore.InMemory.FunctionalTests/Query/GearsOfWarQueryInMemoryTest.cs

@@ -94,4 +94,19 @@ public override Task Null_semantics_is_correctly_applied_for_function_comparison
         // Null protection. Issue #13721.
         => Assert.ThrowsAsync<InvalidOperationException>(
             () => base.Null_semantics_is_correctly_applied_for_function_comparisons_that_take_arguments_from_optional_navigation(async));
+
+    public override Task ElementAt_basic_with_OrderBy(bool async)
+        => Task.CompletedTask;
+
+    public override Task ElementAtOrDefault_basic_with_OrderBy(bool async)
+        => Task.CompletedTask;
+
+    public override Task ElementAtOrDefault_basic_with_OrderBy_parameter(bool async)
+        => Task.CompletedTask;
+
+    public override Task Where_subquery_with_ElementAtOrDefault_equality_to_null_with_composite_key(bool async)
+        => Task.CompletedTask;
+
+    public override Task Where_subquery_with_ElementAt_using_column_as_index(bool async)
+        => Task.CompletedTask;
 }

test/EFCore.InMemory.FunctionalTests/Query/NorthwindMiscellaneousQueryInMemoryTest.cs

@@ -56,4 +56,13 @@ public override async Task Entity_equality_through_subquery_composite_key(bool a
             CoreStrings.EntityEqualityOnCompositeKeyEntitySubqueryNotSupported("==", nameof(OrderDetail)),
             (await Assert.ThrowsAsync<InvalidOperationException>(
                 () => base.Entity_equality_through_subquery_composite_key(async))).Message);
+
+    public override Task Collection_navigation_equal_to_null_for_subquery_using_ElementAtOrDefault_constant_zero(bool async)
+        => Task.CompletedTask;
+
+    public override Task Collection_navigation_equal_to_null_for_subquery_using_ElementAtOrDefault_constant_one(bool async)
+        => Task.CompletedTask;
+
+    public override Task Collection_navigation_equal_to_null_for_subquery_using_ElementAtOrDefault_parameter(bool async)
+        => Task.CompletedTask;
 }

test/EFCore.InMemory.FunctionalTests/Query/NorthwindWhereQueryInMemoryTest.cs

@@ -29,4 +29,10 @@ public override async Task<string> Where_simple_closure(bool async)
     public override Task Like_with_non_string_column_using_double_cast(bool async)
         // Casting int to object to string is invalid for InMemory
         => Assert.ThrowsAsync<InvalidCastException>(() => base.Like_with_non_string_column_using_double_cast(async));
+
+    public override Task ElementAt_over_custom_projection_compared_to_not_null(bool async)
+        => Task.CompletedTask;
+
+    public override Task ElementAtOrDefault_over_custom_projection_compared_to_null(bool async)
+        => Task.CompletedTask;
 }

test/EFCore.Relational.Specification.Tests/Query/ComplexNavigationsCollectionsSplitQueryRelationalTestBase.cs

@@ -13,7 +13,11 @@ protected ComplexNavigationsCollectionsSplitQueryRelationalTestBase(TFixture fix
     }
 
     protected override Expression RewriteServerQueryExpression(Expression serverQueryExpression)
-        => new SplitQueryRewritingExpressionVisitor().Visit(serverQueryExpression);
+    {
+        serverQueryExpression = base.RewriteServerQueryExpression(serverQueryExpression);
+
+        return new SplitQueryRewritingExpressionVisitor().Visit(serverQueryExpression);
+    }
 
     private class SplitQueryRewritingExpressionVisitor : ExpressionVisitor
     {