Query: Update column references in pending collections (#24491)

- Don't apply Include on entities with Include already applied
- Update table references when pushing down select into left for set operation
- Update identifiers after applying set operation if the projection removed exiting identifiers
- Update SQL references in pending collection during push down

Fix for the repro in #17337
Resolves #18738
Resolves #19763
Resolves #19947
Resolves #20813
Resolves #21026
Resolves #22222
Resolves #23676
Resolves #23720
Resolves #24216

src/EFCore.Relational/Query/SqlExpressions/SelectExpression.cs

@@ -1190,11 +1190,23 @@ private void ApplySetOperation(SetOperationType setOperationType, SelectExpressi
             select1._projectionMapping = new Dictionary<ProjectionMember, Expression>(_projectionMapping);
             _projectionMapping.Clear();
 
+            // Remap tableReferences in select1
+            foreach (var tableReference in select1._tableReferences)
+            {
+                tableReference.UpdateTableReference(this, select1);
+            }
+
+            var tableReferenceUpdatingExpressionVisitor = new TableReferenceUpdatingExpressionVisitor(this, select1);
+            tableReferenceUpdatingExpressionVisitor.Visit(select1);
+
             select1._identifier.AddRange(_identifier);
             _identifier.Clear();
             var outerIdentifiers = select1._identifier.Count == select2._identifier.Count
                 ? new ColumnExpression?[select1._identifier.Count]
                 : Array.Empty<ColumnExpression?>();
+            var entityProjectionIdentifiers = new List<ColumnExpression>();
+            var entityProjectionValueComparers = new List<ValueComparer>();
+            var otherExpressions = new List<SqlExpression>();
 
             if (select1.Orderings.Count != 0
                 || select1.Limit != null
@@ -1306,14 +1318,34 @@ private void ApplySetOperation(SetOperationType setOperationType, SelectExpressi
                             outerIdentifiers = Array.Empty<ColumnExpression>();
                         }
                     }
+
+                    otherExpressions.Add(outerProjection);
                 }
             }
 
+            // We should apply _identifiers only when it is distinct and actual select expression had identifiers.
             if (distinct
-                && outerIdentifiers.Length > 0
-                && outerIdentifiers.All(e => e != null))
+                && outerIdentifiers.Length > 0)
             {
-                _identifier.AddRange(outerIdentifiers.Zip(select1._identifier, (c, i) => (c!, i.Comparer)));
+                // If we find matching identifier in outer level then we just use them.
+                if (outerIdentifiers.All(e => e != null))
+                {
+                    _identifier.AddRange(outerIdentifiers.Zip(select1._identifier, (c, i) => (c!, i.Comparer)));
+                }
+                else
+                {
+                    _identifier.Clear();
+                    if (otherExpressions.Count == 0)
+                    {
+                        // If there are no other expressions then we can use all entityProjectionIdentifiers
+                        _identifier.AddRange(entityProjectionIdentifiers.Zip(entityProjectionValueComparers));
+                    }
+                    else if (otherExpressions.All(e => e is ColumnExpression))
+                    {
+                        _identifier.AddRange(entityProjectionIdentifiers.Zip(entityProjectionValueComparers));
+                        _identifier.AddRange(otherExpressions.Select(e => ((ColumnExpression)e, e.TypeMapping!.KeyComparer)));
+                    }
+                }
             }
 
             void HandleEntityProjection(
@@ -1377,8 +1409,22 @@ void HandleEntityProjection(
                     discriminatorExpression = new ConcreteColumnExpression(innerProjection, tableReferenceExpression);
                 }
 
-                _projectionMapping[projectionMember] = new EntityProjectionExpression(
-                    projection1.EntityType, propertyExpressions, discriminatorExpression);
+                var entityProjection = new EntityProjectionExpression(projection1.EntityType, propertyExpressions, discriminatorExpression);
+
+                if (outerIdentifiers.Length > 0)
+                {
+                    var primaryKey = entityProjection.EntityType.FindPrimaryKey();
+                    // If there are any existing identifier then all entity projection must have a key
+                    // else keyless entity would have wiped identifier when generating join.
+                    Check.DebugAssert(primaryKey != null, "primary key is null.");
+                    foreach (var property in primaryKey.Properties)
+                    {
+                        entityProjectionIdentifiers.Add(entityProjection.BindProperty(property));
+                        entityProjectionValueComparers.Add(property.GetKeyValueComparer());
+                    }
+                }
+
+                _projectionMapping[projectionMember] = entityProjection;
             }
 
             string GenerateUniqueColumnAlias(string baseAlias)
@@ -2831,6 +2877,14 @@ protected override Expression VisitChildren(ExpressionVisitor visitor)
                 Offset = (SqlExpression?)visitor.Visit(Offset);
                 Limit = (SqlExpression?)visitor.Visit(Limit);
 
+                if (visitor is SqlRemappingVisitor)
+                {
+                    // We have to traverse pending collections for remapping so that columns from outer are updated.
+                    var pendingCollections = _pendingCollections.ToList();
+                    _pendingCollections.Clear();
+                    _pendingCollections.AddRange(pendingCollections.Select(e => (SelectExpression)visitor.Visit(e)!));
+                }
+
                 return this;
             }
             var changed = false;

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

@@ -538,6 +538,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
                         return ExpandInclude(ownedNavigationReference, ownedNavigationReference.EntityReference);
 
                     case MaterializeCollectionNavigationExpression _:
+                    case IncludeExpression _:
                         return extensionExpression;
                 }
 

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

@@ -1284,6 +1284,24 @@ public override Task Collection_include_over_result_of_single_non_scalar(bool as
             return base.Collection_include_over_result_of_single_non_scalar(async);
         }
 
+        [ConditionalTheory(Skip = "Cross collection join Issue#17246")]
+        public override Task Collection_projection_selecting_outer_element_followed_by_take(bool async)
+        {
+            return base.Collection_projection_selecting_outer_element_followed_by_take(async);
+        }
+
+        [ConditionalTheory(Skip = "Cross collection join Issue#17246")]
+        public override Task Take_on_top_level_and_on_collection_projection_with_outer_apply(bool async)
+        {
+            return base.Take_on_top_level_and_on_collection_projection_with_outer_apply(async);
+        }
+
+        [ConditionalTheory(Skip = "Cross collection join Issue#17246")]
+        public override Task Take_on_correlated_collection_in_first(bool async)
+        {
+            return base.Take_on_correlated_collection_in_first(async);
+        }
+
         private void AssertSql(params string[] expected)
             => Fixture.TestSqlLoggerFactory.AssertBaseline(expected);
 

test/EFCore.Specification.Tests/Query/ComplexNavigationsQueryTestBase.cs

@@ -6093,5 +6093,81 @@ orderby l2.OneToOne_Optional_FK2.MaybeScalar(e => e.Id)
                     AssertEqual(e.Property, a.Property);
                 });
         }
+
+        [ConditionalTheory]
+        [MemberData(nameof(IsAsyncData))]
+        public virtual Task Take_Select_collection_Take(bool async)
+        {
+            return AssertQuery(
+                async,
+                ss => ss.Set<Level1>().OrderBy(l1 => l1.Id).Take(1)
+                    .Select(l1 => new
+                    {
+                        Id = l1.Id,
+                        Name = l1.Name,
+                        Level2s = l1.OneToMany_Required1.OrderBy(l2 => l2.Id).Take(3)
+                            .Select(l2 => new
+                            {
+                                Id = l2.Id,
+                                Name = l2.Name,
+                                Level1Id = EF.Property<int>(l2, "OneToMany_Required_Inverse2Id"),
+                                Level2Id = l2.Level1_Required_Id,
+                                Level2 = l2.OneToOne_Required_FK_Inverse2
+                            })
+                    }),
+                assertOrder: true,
+                elementAsserter: (e, a) =>
+                {
+                    Assert.Equal(e.Id, a.Id);
+                    Assert.Equal(e.Name, a.Name);
+                    AssertCollection(e.Level2s, a.Level2s, ordered: true,
+                        elementAsserter: (ee, aa) =>
+                        {
+                            Assert.Equal(ee.Id, aa.Id);
+                            Assert.Equal(ee.Name, aa.Name);
+                            Assert.Equal(ee.Level1Id, aa.Level1Id);
+                            Assert.Equal(ee.Level2Id, aa.Level2Id);
+                            AssertEqual(ee.Level2, aa.Level2);
+                        });
+                });
+        }
+
+        [ConditionalTheory]
+        [MemberData(nameof(IsAsyncData))]
+        public virtual Task Skip_Take_Select_collection_Skip_Take(bool async)
+        {
+            return AssertQuery(
+                async,
+                ss => ss.Set<Level1>().OrderBy(l1 => l1.Id).Skip(1).Take(1)
+                    .Select(l1 => new
+                    {
+                        Id = l1.Id,
+                        Name = l1.Name,
+                        Level2s = l1.OneToMany_Required1.OrderBy(l2 => l2.Id).Skip(1).Take(3)
+                            .Select(l2 => new
+                            {
+                                Id = l2.Id,
+                                Name = l2.Name,
+                                Level1Id = EF.Property<int>(l2, "OneToMany_Required_Inverse2Id"),
+                                Level2Id = l2.Level1_Required_Id,
+                                Level2 = l2.OneToOne_Required_FK_Inverse2
+                            })
+                    }),
+                assertOrder: true,
+                elementAsserter: (e, a) =>
+                {
+                    Assert.Equal(e.Id, a.Id);
+                    Assert.Equal(e.Name, a.Name);
+                    AssertCollection(e.Level2s, a.Level2s, ordered: true,
+                        elementAsserter: (ee, aa) =>
+                        {
+                            Assert.Equal(ee.Id, aa.Id);
+                            Assert.Equal(ee.Name, aa.Name);
+                            Assert.Equal(ee.Level1Id, aa.Level1Id);
+                            Assert.Equal(ee.Level2Id, aa.Level2Id);
+                            AssertEqual(ee.Level2, aa.Level2);
+                        });
+                });
+        }
     }
 }

test/EFCore.Specification.Tests/Query/NorthwindSelectQueryTestBase.cs

@@ -2168,6 +2168,8 @@ public virtual Task Ternary_in_client_eval_assigns_correct_types(bool async)
               });
         }
 
+        private static string ClientMethod(string s) => s;
+
         [ConditionalTheory]
         [MemberData(nameof(IsAsyncData))]
         public virtual Task VisitLambda_should_not_be_visited_trivially(bool async)
@@ -2246,6 +2248,74 @@ public virtual Task Collection_include_over_result_of_single_non_scalar(bool asy
               entryCount: 235);
         }
 
-        private static string ClientMethod(string s) => s;
+        [ConditionalTheory]
+        [MemberData(nameof(IsAsyncData))]
+        public virtual Task Collection_projection_selecting_outer_element_followed_by_take(bool async)
+        {
+            return AssertQuery(
+              async,
+              ss =>
+               ss.Set<Customer>().Include(c => c.Orders)
+                    .Where(c => c.CustomerID.StartsWith("F"))
+                    .OrderBy(e => e.CustomerID)
+                    .Select(c => new
+                    {
+                        Customer = c.Orders.Select(o => c)
+                    })
+                    .Take(10),
+              assertOrder: true,
+              elementAsserter: (e, a) =>
+              {
+                  AssertCollection(e.Customer, a.Customer,
+                      elementAsserter: (ee, aa) => AssertInclude(ee, aa, new ExpectedInclude<Customer>(i => i.Orders)));
+              },
+              entryCount: 70);
+        }
+
+        [ConditionalTheory]
+        [MemberData(nameof(IsAsyncData))]
+        public virtual Task Take_on_top_level_and_on_collection_projection_with_outer_apply(bool async)
+        {
+            return AssertFirstOrDefault(
+              async,
+              ss =>
+               ss.Set<Order>()
+                    .Where(o => o.CustomerID.StartsWith("F"))
+                    .Select(o => new Order
+                    {
+                        OrderID = o.OrderID,
+                        OrderDate = o.OrderDate,
+                        OrderDetails = o.OrderDetails.Select(e => new OrderDetail
+                        {
+                            OrderID = e.OrderID,
+                            Product = e.Product,
+                            UnitPrice = e.UnitPrice
+                        })
+                        .OrderByDescending(e => e.OrderID)
+                        .Skip(0)
+                        .Take(10)
+                        .ToList()
+                    }),
+              entryCount: 2);
+        }
+
+        [ConditionalTheory]
+        [MemberData(nameof(IsAsyncData))]
+        public virtual Task Take_on_correlated_collection_in_first(bool async)
+        {
+            return AssertFirstOrDefault(
+              async,
+              ss =>
+               ss.Set<Customer>()
+                    .Where(o => o.CustomerID.StartsWith("F"))
+                    .OrderBy(e => e.CustomerID)
+                    .Select(o => new
+                    {
+                        Orders = o.Orders.OrderBy(a => a.OrderDate).Take(1)
+                            .Select(e => new { Title = e.CustomerID == e.Customer.CustomerID ? "A" : "B" }).ToList()
+                    }),
+              asserter: (e, a) => AssertCollection(e.Orders, a.Orders, ordered: true,
+                elementAsserter: (ee, aa) => AssertEqual(ee.Title, aa.Title)));
+        }
     }
 }