Prune PgTableValuedFunctionExpression's projection (#3090)

Closes #3088

src/EFCore.PG/Extensions/ExpressionVisitorExtensions.cs

@@ -16,20 +16,21 @@ internal static class ExpressionVisitorExtensions
     /// <returns>
     ///     The modified expression list, if any of the elements were modified; otherwise, returns the original expression list.
     /// </returns>
-    public static IReadOnlyList<Expression> Visit(this ExpressionVisitor visitor, IReadOnlyList<Expression> nodes)
+    public static IReadOnlyList<T> Visit<T>(this ExpressionVisitor visitor, IReadOnlyList<T> nodes)
+        where T : Expression
     {
-        Expression[]? newNodes = null;
+        T[]? newNodes = null;
         for (int i = 0, n = nodes.Count; i < n; i++)
         {
-            var node = visitor.Visit(nodes[i]);
+            var node = (T)visitor.Visit(nodes[i]);
 
             if (newNodes is not null)
             {
                 newNodes[i] = node;
             }
             else if (!ReferenceEquals(node, nodes[i]))
             {
-                newNodes = new Expression[n];
+                newNodes = new T[n];
                 for (var j = 0; j < i; j++)
                 {
                     newNodes[j] = nodes[j];

src/EFCore.PG/Query/Expressions/Internal/PgTableValuedFunctionExpression.cs

@@ -96,6 +96,15 @@ public override TableExpressionBase Clone(string? alias, ExpressionVisitor cloni
     public override PgTableValuedFunctionExpression WithAlias(string newAlias)
         => new(newAlias, Name, Arguments, ColumnInfos, WithOrdinality);
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public virtual PgTableValuedFunctionExpression WithColumnInfos(IReadOnlyList<ColumnInfo> columnInfos)
+        => new(Alias, Name, Arguments, columnInfos, WithOrdinality);
+
     /// <inheritdoc />
     protected override void Print(ExpressionPrinter expressionPrinter)
     {

src/EFCore.PG/Query/Expressions/Internal/PgUnnestExpression.cs

@@ -54,7 +54,12 @@ public virtual string ColumnName
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public PgUnnestExpression(string alias, SqlExpression array, string columnName, bool withOrdinality = true)
-        : base(alias, "unnest", new[] { array }, new[] { new ColumnInfo(columnName) }, withOrdinality)
+        : this(alias, array, new ColumnInfo(columnName), withOrdinality)
+    {
+    }
+
+    private PgUnnestExpression(string alias, SqlExpression array, ColumnInfo? columnInfo, bool withOrdinality = true)
+        : base(alias, "unnest", new[] { array }, columnInfo is null ? null : [columnInfo.Value], withOrdinality)
     {
     }
 
@@ -91,6 +96,19 @@ public override TableExpressionBase Clone(string? alias, ExpressionVisitor cloni
         => new PgUnnestExpression(alias!, (SqlExpression)cloningExpressionVisitor.Visit(Array), ColumnName, WithOrdinality);
 
     /// <inheritdoc />
-    public override PgTableValuedFunctionExpression WithAlias(string newAlias)
-        => new(newAlias, Name, Arguments, ColumnInfos, WithOrdinality);
+    public override PgUnnestExpression WithAlias(string newAlias)
+        => new(newAlias, Array, ColumnName, WithOrdinality);
+
+    /// <inheritdoc />
+    public override PgUnnestExpression WithColumnInfos(IReadOnlyList<ColumnInfo> columnInfos)
+        => new(
+            Alias,
+            Array,
+            columnInfos switch
+            {
+                [] => null,
+                [var columnInfo] => columnInfo,
+                _ => throw new ArgumentException()
+            },
+            WithOrdinality);
 }

src/EFCore.PG/Query/Internal/NpgsqlQueryTranslationPostprocessor.cs

@@ -8,6 +8,8 @@ namespace Npgsql.EntityFrameworkCore.PostgreSQL.Query.Internal;
 /// </summary>
 public class NpgsqlQueryTranslationPostprocessor : RelationalQueryTranslationPostprocessor
 {
+    private readonly NpgsqlSqlTreePruner _pruner = new();
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -37,4 +39,13 @@ public override Expression Process(Expression query)
 
         return result;
     }
+
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    protected override Expression Prune(Expression query)
+        => _pruner.Prune(query);
 }

src/EFCore.PG/Query/Internal/NpgsqlSqlTreePruner.cs

@@ -0,0 +1,56 @@
+using Npgsql.EntityFrameworkCore.PostgreSQL.Query.Expressions.Internal;
+using ColumnInfo = Npgsql.EntityFrameworkCore.PostgreSQL.Query.Expressions.Internal.PgTableValuedFunctionExpression.ColumnInfo;
+
+namespace Npgsql.EntityFrameworkCore.PostgreSQL.Query.Internal;
+
+/// <summary>
+///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+///     any release. You should only use it directly in your code with extreme caution and knowing that
+///     doing so can result in application failures when updating to a new Entity Framework Core release.
+/// </summary>
+public class NpgsqlSqlTreePruner  : SqlTreePruner
+{
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    protected override Expression VisitExtension(Expression node)
+    {
+        switch (node)
+        {
+            case PgTableValuedFunctionExpression { ColumnInfos: IReadOnlyList<ColumnInfo> columnInfos } tvf:
+                var arguments = this.Visit(tvf.Arguments);
+
+                List<ColumnInfo>? newColumnInfos = null;
+
+                if (ReferencedColumnMap.TryGetValue(tvf.Alias, out var referencedAliases))
+                {
+                    for (var i = 0; i < columnInfos.Count; i++)
+                    {
+                        if (referencedAliases.Contains(columnInfos[i].Name))
+                        {
+                            newColumnInfos?.Add(columnInfos[i]);
+                        }
+                        else if (newColumnInfos is null)
+                        {
+                            newColumnInfos = [];
+                            for (var j = 0; j < i; j++)
+                            {
+                                newColumnInfos.Add(columnInfos[j]);
+                            }
+                        }
+                    }
+                }
+
+                return tvf
+                    .Update(arguments)
+                    .WithColumnInfos(newColumnInfos ?? columnInfos);
+
+            default:
+                return base.VisitExtension(node);
+        }
+    }
+}

src/EFCore.PG/Query/Internal/NpgsqlUnnestPostprocessor.cs

@@ -36,16 +36,16 @@ public class NpgsqlUnnestPostprocessor : ExpressionVisitor
                 for (var i = 0; i < selectExpression.Tables.Count; i++)
                 {
                     var table = selectExpression.Tables[i];
+                    var unwrappedTable = table.UnwrapJoin();
 
                     // Find any unnest table which does not have any references to its ordinality column in the projection or orderings
                     // (this is where they may appear when a column is an identifier).
-                    var unnest = table as PgUnnestExpression ?? (table as JoinExpressionBase)?.Table as PgUnnestExpression;
-                    if (unnest is not null
+                    if (unwrappedTable is PgUnnestExpression unnest
                         && !selectExpression.Orderings.Select(o => o.Expression)
                             .Concat(selectExpression.Projection.Select(p => p.Expression))
                             .Any(
-                                p => p is ColumnExpression { Name: "ordinality", } ordinalityColumn
-                                    && ordinalityColumn.TableAlias == table.Alias))
+                                p => p is ColumnExpression { Name: "ordinality" } ordinalityColumn
+                                    && ordinalityColumn.TableAlias == unwrappedTable.Alias))
                     {
                         if (newTables is null)
                         {

test/EFCore.PG.FunctionalTests/Query/JsonQueryNpgsqlTest.cs

@@ -1032,16 +1032,7 @@ public override async Task Json_collection_Any_with_predicate(bool async)
 FROM "JsonEntitiesBasic" AS j
 WHERE EXISTS (
     SELECT 1
-    FROM ROWS FROM (jsonb_to_recordset(j."OwnedReferenceRoot" -> 'OwnedCollectionBranch') AS (
-        "Date" timestamp without time zone,
-        "Enum" integer,
-        "Enums" integer[],
-        "Fraction" numeric(18,2),
-        "NullableEnum" integer,
-        "NullableEnums" integer[],
-        "OwnedCollectionLeaf" jsonb,
-        "OwnedReferenceLeaf" jsonb
-    )) WITH ORDINALITY AS o
+    FROM ROWS FROM (jsonb_to_recordset(j."OwnedReferenceRoot" -> 'OwnedCollectionBranch') AS ("OwnedReferenceLeaf" jsonb)) WITH ORDINALITY AS o
     WHERE (o."OwnedReferenceLeaf" ->> 'SomethingSomething') = 'e1_r_c1_r')
 """);
     }
@@ -1059,11 +1050,7 @@ SELECT o."OwnedReferenceLeaf" ->> 'SomethingSomething'
     FROM ROWS FROM (jsonb_to_recordset(j."OwnedReferenceRoot" -> 'OwnedCollectionBranch') AS (
         "Date" timestamp without time zone,
         "Enum" integer,
-        "Enums" integer[],
         "Fraction" numeric(18,2),
-        "NullableEnum" integer,
-        "NullableEnums" integer[],
-        "OwnedCollectionLeaf" jsonb,
         "OwnedReferenceLeaf" jsonb
     )) WITH ORDINALITY AS o
     WHERE o."Enum" = -3
@@ -1083,16 +1070,7 @@ public override async Task Json_collection_Skip(bool async)
     SELECT o0.c
     FROM (
         SELECT o."OwnedReferenceLeaf" ->> 'SomethingSomething' AS c
-        FROM ROWS FROM (jsonb_to_recordset(j."OwnedReferenceRoot" -> 'OwnedCollectionBranch') AS (
-            "Date" timestamp without time zone,
-            "Enum" integer,
-            "Enums" integer[],
-            "Fraction" numeric(18,2),
-            "NullableEnum" integer,
-            "NullableEnums" integer[],
-            "OwnedCollectionLeaf" jsonb,
-            "OwnedReferenceLeaf" jsonb
-        )) WITH ORDINALITY AS o
+        FROM ROWS FROM (jsonb_to_recordset(j."OwnedReferenceRoot" -> 'OwnedCollectionBranch') AS ("OwnedReferenceLeaf" jsonb)) WITH ORDINALITY AS o
         OFFSET 1
     ) AS o0
     LIMIT 1 OFFSET 0) = 'e1_r_c2_r'
@@ -1114,11 +1092,7 @@ SELECT o0.c
         FROM ROWS FROM (jsonb_to_recordset(j."OwnedReferenceRoot" -> 'OwnedCollectionBranch') AS (
             "Date" timestamp without time zone,
             "Enum" integer,
-            "Enums" integer[],
             "Fraction" numeric(18,2),
-            "NullableEnum" integer,
-            "NullableEnums" integer[],
-            "OwnedCollectionLeaf" jsonb,
             "OwnedReferenceLeaf" jsonb
         )) WITH ORDINALITY AS o
         ORDER BY o."Date" DESC NULLS LAST
@@ -1166,14 +1140,7 @@ public override async Task Json_collection_within_collection_Count(bool async)
 FROM "JsonEntitiesBasic" AS j
 WHERE EXISTS (
     SELECT 1
-    FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS (
-        "Name" text,
-        "Names" text[],
-        "Number" integer,
-        "Numbers" integer[],
-        "OwnedCollectionBranch" jsonb,
-        "OwnedReferenceBranch" jsonb
-    )) WITH ORDINALITY AS o
+    FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS ("OwnedCollectionBranch" jsonb)) WITH ORDINALITY AS o
     WHERE (
         SELECT count(*)::int
         FROM ROWS FROM (jsonb_to_recordset(o."OwnedCollectionBranch") AS (
@@ -1220,11 +1187,7 @@ public override async Task Json_collection_in_projection_with_anonymous_projecti
 FROM "JsonEntitiesBasic" AS j
 LEFT JOIN LATERAL ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS (
     "Name" text,
-    "Names" text[],
-    "Number" integer,
-    "Numbers" integer[],
-    "OwnedCollectionBranch" jsonb,
-    "OwnedReferenceBranch" jsonb
+    "Number" integer
 )) WITH ORDINALITY AS o ON TRUE
 ORDER BY j."Id" NULLS FIRST
 """);
@@ -1242,11 +1205,7 @@ LEFT JOIN LATERAL (
     SELECT o."Name", o."Number", o.ordinality
     FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS (
         "Name" text,
-        "Names" text[],
-        "Number" integer,
-        "Numbers" integer[],
-        "OwnedCollectionBranch" jsonb,
-        "OwnedReferenceBranch" jsonb
+        "Number" integer
     )) WITH ORDINALITY AS o
     WHERE o."Name" = 'Foo'
 ) AS o0 ON TRUE
@@ -1268,9 +1227,7 @@ FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS (
         "Name" text,
         "Names" text[],
         "Number" integer,
-        "Numbers" integer[],
-        "OwnedCollectionBranch" jsonb,
-        "OwnedReferenceBranch" jsonb
+        "Numbers" integer[]
     )) WITH ORDINALITY AS o
     WHERE o."Name" = 'Foo'
 ) AS o0 ON TRUE
@@ -1312,14 +1269,7 @@ public override async Task Json_nested_collection_filter_in_projection(bool asyn
 FROM "JsonEntitiesBasic" AS j
 LEFT JOIN LATERAL (
     SELECT o.ordinality, o1."Id", o1."Date", o1."Enum", o1."Enums", o1."Fraction", o1."NullableEnum", o1."NullableEnums", o1.c, o1.c0, o1.ordinality AS ordinality0
-    FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS (
-        "Name" text,
-        "Names" text[],
-        "Number" integer,
-        "Numbers" integer[],
-        "OwnedCollectionBranch" jsonb,
-        "OwnedReferenceBranch" jsonb
-    )) WITH ORDINALITY AS o
+    FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS ("OwnedCollectionBranch" jsonb)) WITH ORDINALITY AS o
     LEFT JOIN LATERAL (
         SELECT j."Id", o0."Date", o0."Enum", o0."Enums", o0."Fraction", o0."NullableEnum", o0."NullableEnums", o0."OwnedCollectionLeaf" AS c, o0."OwnedReferenceLeaf" AS c0, o0.ordinality
         FROM ROWS FROM (jsonb_to_recordset(o."OwnedCollectionBranch") AS (
@@ -1349,21 +1299,12 @@ public override async Task Json_nested_collection_anonymous_projection_in_projec
 FROM "JsonEntitiesBasic" AS j
 LEFT JOIN LATERAL (
     SELECT o.ordinality, o0."Date" AS c, o0."Enum" AS c0, o0."Enums" AS c1, o0."Fraction" AS c2, o0."OwnedReferenceLeaf" AS c3, j."Id", o0."OwnedCollectionLeaf" AS c4, o0.ordinality AS ordinality0
-    FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS (
-        "Name" text,
-        "Names" text[],
-        "Number" integer,
-        "Numbers" integer[],
-        "OwnedCollectionBranch" jsonb,
-        "OwnedReferenceBranch" jsonb
-    )) WITH ORDINALITY AS o
+    FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS ("OwnedCollectionBranch" jsonb)) WITH ORDINALITY AS o
     LEFT JOIN LATERAL ROWS FROM (jsonb_to_recordset(o."OwnedCollectionBranch") AS (
         "Date" timestamp without time zone,
         "Enum" integer,
         "Enums" integer[],
         "Fraction" numeric(18,2),
-        "NullableEnum" integer,
-        "NullableEnums" integer[],
         "OwnedCollectionLeaf" jsonb,
         "OwnedReferenceLeaf" jsonb
     )) WITH ORDINALITY AS o0 ON TRUE
@@ -1434,10 +1375,7 @@ LEFT JOIN LATERAL (
     SELECT o."OwnedReferenceBranch" AS c, j."Id", o.ordinality, o."Name" AS c0
     FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS (
         "Name" text,
-        "Names" text[],
         "Number" integer,
-        "Numbers" integer[],
-        "OwnedCollectionBranch" jsonb,
         "OwnedReferenceBranch" jsonb
     )) WITH ORDINALITY AS o
     ORDER BY o."Name" NULLS FIRST
@@ -1520,14 +1458,7 @@ FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS (
 ) AS o1 ON TRUE
 LEFT JOIN LATERAL (
     SELECT o2.ordinality, o5."Id", o5."Date", o5."Enum", o5."Enums", o5."Fraction", o5."NullableEnum", o5."NullableEnums", o5.c, o5.c0, o5.ordinality AS ordinality0
-    FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS (
-        "Name" text,
-        "Names" text[],
-        "Number" integer,
-        "Numbers" integer[],
-        "OwnedCollectionBranch" jsonb,
-        "OwnedReferenceBranch" jsonb
-    )) WITH ORDINALITY AS o2
+    FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS ("OwnedCollectionBranch" jsonb)) WITH ORDINALITY AS o2
     LEFT JOIN LATERAL (
         SELECT j."Id", o3."Date", o3."Enum", o3."Enums", o3."Fraction", o3."NullableEnum", o3."NullableEnums", o3."OwnedCollectionLeaf" AS c, o3."OwnedReferenceLeaf" AS c0, o3.ordinality
         FROM ROWS FROM (jsonb_to_recordset(o2."OwnedCollectionBranch") AS (
@@ -1756,14 +1687,7 @@ public override async Task Json_collection_Select_entity_in_anonymous_object_Ele
 FROM "JsonEntitiesBasic" AS j
 LEFT JOIN LATERAL (
     SELECT o."OwnedReferenceBranch" AS c, j."Id", 1 AS c0
-    FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS (
-        "Name" text,
-        "Names" text[],
-        "Number" integer,
-        "Numbers" integer[],
-        "OwnedCollectionBranch" jsonb,
-        "OwnedReferenceBranch" jsonb
-    )) WITH ORDINALITY AS o
+    FROM ROWS FROM (jsonb_to_recordset(j."OwnedCollectionRoot") AS ("OwnedReferenceBranch" jsonb)) WITH ORDINALITY AS o
     LIMIT 1 OFFSET 0
 ) AS o0 ON TRUE
 ORDER BY j."Id" NULLS FIRST

test/EFCore.PG.FunctionalTests/Query/PrimitiveCollectionsQueryNpgsqlTest.cs

@@ -1312,7 +1312,7 @@ public override async Task Project_collection_of_ints_with_distinct(bool async)
 FROM "PrimitiveCollectionsEntity" AS p
 LEFT JOIN LATERAL (
     SELECT DISTINCT i.value
-    FROM unnest(p."Ints") WITH ORDINALITY AS i(value)
+    FROM unnest(p."Ints") AS i(value)
 ) AS i0 ON TRUE
 ORDER BY p."Id" NULLS FIRST
 """);