Bucketization for IN (#36370)

Author: cincuranet

src/EFCore.Relational/Query/SqlNullabilityProcessor.cs

@@ -131,23 +131,18 @@ protected override Expression VisitExtension(Expression node)
                     case ParameterTranslationMode.MultipleParameters:
                     {
                         var expandedParameters = _collectionParameterExpansionMap.GetOrAddNew(valuesParameter);
-                        for (var i = 0; i < values.Count; i++)
+                        var expandedParametersCounter = 0;
+                        foreach (var value in values)
                         {
                             // Create parameter for value if we didn't create it yet,
                             // otherwise reuse it.
-                            if (expandedParameters.Count <= i)
-                            {
-                                var parameterName = Uniquifier.Uniquify(valuesParameter.Name, queryParameters, int.MaxValue);
-                                queryParameters.Add(parameterName, values[i]);
-                                var parameterExpression = new SqlParameterExpression(parameterName, values[i]?.GetType() ?? typeof(object), elementTypeMapping);
-                                expandedParameters.Add(parameterExpression);
-                            }
+                            ExpandParameterIfNeeded(valuesParameter.Name, expandedParameters, queryParameters, expandedParametersCounter, value, elementTypeMapping);
 
                             processedValues.Add(
                                 new RowValueExpression(
                                     ProcessValuesOrderingColumn(
                                         valuesExpression,
-                                        [expandedParameters[i]],
+                                        [expandedParameters[expandedParametersCounter++]],
                                         intTypeMapping,
                                         ref valuesOrderingCounter)));
                         }
@@ -814,41 +809,35 @@ InExpression ProcessInExpressionValues(
 
                 processedValues = [];
 
+                var translationMode = valuesParameter.TranslationMode ?? CollectionParameterTranslationMode;
                 var expandedParameters = _collectionParameterExpansionMap.GetOrAddNew(valuesParameter);
                 var expandedParametersCounter = 0;
-                for (var i = 0; i < values.Count; i++)
+                foreach (var value in values)
                 {
-                    if (values[i] is null && removeNulls)
+                    if (value is null && removeNulls)
                     {
                         hasNull = true;
                         continue;
                     }
 
-                    switch (valuesParameter.TranslationMode ?? CollectionParameterTranslationMode)
+                    switch (translationMode)
                     {
                         case ParameterTranslationMode.MultipleParameters:
                         // see #36311 for more info
                         case ParameterTranslationMode.Parameter:
                         {
                             // Create parameter for value if we didn't create it yet,
                             // otherwise reuse it.
-                            if (expandedParameters.Count <= i)
-                            {
-                                var parameterName = Uniquifier.Uniquify(valuesParameter.Name, parameters, int.MaxValue);
-                                parameters.Add(parameterName, values[i]);
-                                var parameterExpression = new SqlParameterExpression(parameterName, values[i]?.GetType() ?? typeof(object), elementTypeMapping);
-                                expandedParameters.Add(parameterExpression);
-                            }
+                            ExpandParameterIfNeeded(valuesParameter.Name, expandedParameters, parameters, expandedParametersCounter, value, elementTypeMapping);
 
-                            // Use separate counter, because we may skip nulls.
                             processedValues.Add(expandedParameters[expandedParametersCounter++]);
 
                             break;
                         }
 
                         case ParameterTranslationMode.Constant:
                         {
-                            processedValues.Add(_sqlExpressionFactory.Constant(values[i], values[i]?.GetType() ?? typeof(object), sensitive: true, elementTypeMapping));
+                            processedValues.Add(_sqlExpressionFactory.Constant(value, value?.GetType() ?? typeof(object), sensitive: true, elementTypeMapping));
 
                             break;
                         }
@@ -857,6 +846,38 @@ InExpression ProcessInExpressionValues(
                             throw new UnreachableException();
                     }
                 }
+
+                // Bucketization is a process used to group parameters into "buckets" of a fixed size when generating parameterized collections.
+                // This helps mitigate query plan bloat by reducing the number of unique query plans generated for queries with varying numbers
+                // of parameters. Instead of creating a new query plan for every possible parameter count, bucketization ensures that queries
+                // with similar parameter counts share the same query plan.
+                //
+                // The size of each bucket is determined by the CalculateParameterBucketSize method, which dynamically calculates the bucket size
+                // based on the total number of parameters and the type mapping of the collection elements. For example, smaller collections may
+                // use smaller bucket sizes, while larger collections may use larger bucket sizes to balance performance and memory usage.
+                //
+                // If the number of parameters in the collection is not a multiple of the bucket size, padding is added to ensure the collection
+                // fits into the nearest bucket. This padding uses the last value in the collection to fill the remaining slots.
+                //
+                // Providers can effectively disable bucketization by overriding the CalculateParameterBucketSize method to always return 1.
+                //
+                // Example:
+                // Suppose a query has 12 parameters, and the bucket size is calculated as 10. The query will be padded with 8 additional
+                // parameters (using the last value) to fit into the next bucket size of 20. This ensures that queries with 12, 13, or 19
+                // parameters all share the same query plan, reducing query plan fragmentation.
+                if (translationMode is ParameterTranslationMode.MultipleParameters)
+                {
+                    var padFactor = CalculateParameterBucketSize(values.Count, elementTypeMapping);
+                    var padding = (padFactor - (values.Count % padFactor)) % padFactor;
+                    for (var i = 0; i < padding; i++)
+                    {
+                        // Create parameter for value if we didn't create it yet,
+                        // otherwise reuse it.
+                        ExpandParameterIfNeeded(valuesParameter.Name, expandedParameters, parameters, values.Count + i, values[^1], elementTypeMapping);
+
+                        processedValues.Add(expandedParameters[expandedParametersCounter++]);
+                    }
+                }
             }
             else
             {
@@ -1488,6 +1509,23 @@ protected virtual SqlExpression VisitJsonScalar(
     protected virtual bool PreferExistsToInWithCoalesce
         => false;
 
+    /// <summary>
+    ///     Gets the bucket size into which the parameters are padded when generating a parameterized collection
+    ///     when using multiple parameters. This helps with query plan bloat.
+    /// </summary>
+    /// <param name="count">Number of value parameters.</param>
+    /// <param name="elementTypeMapping">The type mapping for the collection element.</param>
+    [EntityFrameworkInternal]
+    protected virtual int CalculateParameterBucketSize(int count, RelationalTypeMapping elementTypeMapping)
+        => count switch
+        {
+            <= 5 => 1,
+            <= 150 => 10,
+            <= 750 => 50,
+            <= 2000 => 100,
+            _ => 200,
+        };
+
     // Note that we can check parameter values for null since we cache by the parameter nullability; but we cannot do the same for bool.
     private bool IsNull(SqlExpression? expression)
         => expression is SqlConstantExpression { Value: null }
@@ -2121,4 +2159,21 @@ private SqlExpression ProcessNullNotNull(SqlExpression sqlExpression, bool opera
 
     private static bool IsLogicalNot(SqlUnaryExpression? sqlUnaryExpression)
         => sqlUnaryExpression is { OperatorType: ExpressionType.Not } && sqlUnaryExpression.Type == typeof(bool);
+
+    private static void ExpandParameterIfNeeded(
+        string valuesParameterName,
+        List<SqlParameterExpression> expandedParameters,
+        Dictionary<string, object?> parameters,
+        int index,
+        object? value,
+        RelationalTypeMapping typeMapping)
+    {
+        if (expandedParameters.Count <= index)
+        {
+            var parameterName = Uniquifier.Uniquify(valuesParameterName, parameters, int.MaxValue);
+            parameters.Add(parameterName, value);
+            var parameterExpression = new SqlParameterExpression(parameterName, value?.GetType() ?? typeof(object), typeMapping);
+            expandedParameters.Add(parameterExpression);
+        }
+    }
 }

src/EFCore.SqlServer/Query/Internal/SqlServerSqlNullabilityProcessor.cs

@@ -16,6 +16,8 @@ namespace Microsoft.EntityFrameworkCore.SqlServer.Query.Internal;
 /// </summary>
 public class SqlServerSqlNullabilityProcessor : SqlNullabilityProcessor
 {
+    private const int MaxParameterCount = 2100;
+
     /// <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
@@ -280,6 +282,19 @@ protected override SqlExpression VisitIn(InExpression inExpression, bool allowOp
         }
     }
 
+    /// <inheritdoc />
+    protected override int CalculateParameterBucketSize(int count, RelationalTypeMapping elementTypeMapping)
+         => count switch
+         {
+             <= 5 => 1,
+             <= 150 => 10,
+             <= 750 => 50,
+             <= 2000 => 100,
+             <= 2070 => 10, // try not to over-pad as we approach that limit
+             <= MaxParameterCount => 0, // just don't pad between 2070 and 2100, to minimize the crazy
+             _ => 200,
+         };
+
     private bool TryHandleOverLimitParameters(
         SqlParameterExpression valuesParameter,
         RelationalTypeMapping typeMapping,
@@ -294,7 +309,7 @@ private bool TryHandleOverLimitParameters(
         // SQL Server has limit on number of parameters in a query.
         // If we're over that limit, we switch to using single parameter
         // and processing it through JSON functions.
-        if (values.Count > 2098)
+        if (values.Count > MaxParameterCount)
         {
             if (_sqlServerSingletonOptions.SupportsJsonFunctions)
             {

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

@@ -215,6 +215,27 @@ public virtual async Task Parameter_collection_Contains_with_default_mode_EF_Mul
         Assert.Equivalent(new[] { 2 }, result);
     }
 
+    [ConditionalFact]
+    public virtual async Task Parameter_collection_Contains_parameter_bucketization()
+    {
+        var contextFactory = await InitializeAsync<TestContext>(
+            onConfiguring: b => SetParameterizedCollectionMode(b, ParameterTranslationMode.MultipleParameters),
+            seed: context =>
+            {
+                context.AddRange(
+                    new TestEntity { Id = 1 },
+                    new TestEntity { Id = 2 },
+                    new TestEntity { Id = 100 });
+                return context.SaveChangesAsync();
+            });
+
+        await using var context = contextFactory.CreateContext();
+
+        var ints = new[] { 2, 999, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2 };
+        var result = await context.Set<TestEntity>().Where(c => ints.Contains(c.Id)).Select(c => c.Id).ToListAsync();
+        Assert.Equivalent(new[] { 2 }, result);
+    }
+
     protected class TestOwner
     {
         public int Id { get; set; }