[release/10.0] Fallback to OPENJSON with many parameters (#37320)

Author: cincuranet

src/EFCore.Relational/Query/SqlNullabilityProcessor.cs

@@ -926,7 +926,7 @@ InExpression ProcessInExpressionValues(
                 if (translationMode is ParameterTranslationMode.MultipleParameters)
                 {
                     var padFactor = CalculateParameterBucketSize(values.Count, elementTypeMapping);
-                    var padding = (padFactor - (values.Count % padFactor)) % padFactor;
+                    var padding = CalculatePadding(values.Count, padFactor);
                     for (var i = 0; i < padding; i++)
                     {
                         // Create parameter for value if we didn't create it yet,
@@ -1587,6 +1587,15 @@ protected virtual int CalculateParameterBucketSize(int count, RelationalTypeMapp
             _ => 200,
         };
 
+    /// <summary>
+    /// Calculates the number of padding parameters needed to align the total count to the nearest bucket size.
+    /// </summary>
+    /// <param name="count">Number of value parameters.</param>
+    /// <param name="padFactor">Padding factor.</param>
+    [EntityFrameworkInternal]
+    protected virtual int CalculatePadding(int count, int padFactor)
+        => (padFactor - (count % padFactor)) % padFactor;
+
     // 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 }

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

@@ -21,6 +21,9 @@ public class SqlServerSqlNullabilityProcessor : SqlNullabilityProcessor
     private static readonly bool UseOldBehavior37151 =
         AppContext.TryGetSwitch("Microsoft.EntityFrameworkCore.Issue37151", out var enabled) && enabled;
 
+    private static readonly bool UseOldBehavior37185 =
+        AppContext.TryGetSwitch("Microsoft.EntityFrameworkCore.Issue37185", out var enabled) && enabled;
+
     /// <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
@@ -33,6 +36,7 @@ public class SqlServerSqlNullabilityProcessor : SqlNullabilityProcessor
     private readonly ISqlServerSingletonOptions _sqlServerSingletonOptions;
 
     private int _openJsonAliasCounter;
+    private int _totalParameterCount;
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -55,6 +59,18 @@ public SqlServerSqlNullabilityProcessor(
     /// </summary>
     public override Expression Process(Expression queryExpression, ParametersCacheDecorator parametersDecorator)
     {
+        if (!UseOldBehavior37185)
+        {
+            var parametersCounter = new ParametersCounter(
+                parametersDecorator,
+                CollectionParameterTranslationMode,
+#pragma warning disable EF1001
+                (count, elementTypeMapping) => CalculatePadding(count, CalculateParameterBucketSize(count, elementTypeMapping)));
+#pragma warning restore EF1001
+            parametersCounter.Visit(queryExpression);
+            _totalParameterCount = parametersCounter.Count;
+        }
+
         var result = base.Process(queryExpression, parametersDecorator);
         _openJsonAliasCounter = 0;
         return result;
@@ -313,7 +329,9 @@ 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 > MaxParameterCount)
+        if (UseOldBehavior37185
+            ? values.Count > MaxParameterCount
+            : _totalParameterCount > MaxParameterCount)
         {
             if (_sqlServerSingletonOptions.SupportsJsonFunctions)
             {
@@ -372,3 +390,107 @@ valuesExpression is not null
     }
 #pragma warning restore EF1001
 }
+
+/// <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 ParametersCounter(
+    ParametersCacheDecorator parametersDecorator,
+    ParameterTranslationMode collectionParameterTranslationMode,
+    Func<int, RelationalTypeMapping, int> bucketizationPadding) : ExpressionVisitor
+{
+    /// <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 int Count { get; private set; }
+
+    private readonly HashSet<SqlParameterExpression> _visitedSqlParameters =
+        new(EqualityComparer<SqlParameterExpression>.Create(
+            (lhs, rhs) =>
+                ReferenceEquals(lhs, rhs)
+                || (lhs is not null && rhs is not null
+                    && lhs.InvariantName == rhs.InvariantName
+                    && lhs.TranslationMode == rhs.TranslationMode),
+            x => HashCode.Combine(x.InvariantName, x.TranslationMode)));
+
+    private readonly HashSet<QueryParameterExpression> _visitedQueryParameters =
+        new(EqualityComparer<QueryParameterExpression>.Create(
+            (lhs, rhs) =>
+                ReferenceEquals(lhs, rhs)
+                || (lhs is not null && rhs is not null
+                    && lhs.Name == rhs.Name
+                    && lhs.TranslationMode == rhs.TranslationMode),
+            x => HashCode.Combine(x.Name, x.TranslationMode)));
+
+    /// <inheritdoc/>
+    protected override Expression VisitExtension(Expression node)
+    {
+        switch (node)
+        {
+            case ValuesExpression { ValuesParameter: { } valuesParameter }:
+                ProcessCollectionParameter(valuesParameter, bucketization: false);
+                break;
+
+            case InExpression { ValuesParameter: { } valuesParameter }:
+                ProcessCollectionParameter(valuesParameter, bucketization: true);
+                break;
+
+            case FromSqlExpression { Arguments: QueryParameterExpression queryParameter }:
+                if (_visitedQueryParameters.Add(queryParameter))
+                {
+                    var parameters = parametersDecorator.GetAndDisableCaching();
+                    Count += ((object?[])parameters[queryParameter.Name]!).Length;
+                }
+                break;
+
+            case SqlParameterExpression sqlParameterExpression:
+                if (_visitedSqlParameters.Add(sqlParameterExpression))
+                {
+                    Count++;
+                }
+                break;
+        }
+
+        return base.VisitExtension(node);
+    }
+
+    private void ProcessCollectionParameter(SqlParameterExpression sqlParameterExpression, bool bucketization)
+    {
+        if (!_visitedSqlParameters.Add(sqlParameterExpression))
+        {
+            return;
+        }
+
+        switch (sqlParameterExpression.TranslationMode ?? collectionParameterTranslationMode)
+        {
+            case ParameterTranslationMode.MultipleParameters:
+                var parameters = parametersDecorator.GetAndDisableCaching();
+                var count = ((IEnumerable?)parameters[sqlParameterExpression.Name])?.Cast<object?>().Count() ?? 0;
+                Count += count;
+
+                if (bucketization)
+                {
+                    var elementTypeMapping = (RelationalTypeMapping)sqlParameterExpression.TypeMapping!.ElementTypeMapping!;
+                    Count += bucketizationPadding(count, elementTypeMapping);
+                }
+
+                break;
+
+            case ParameterTranslationMode.Parameter:
+                Count++;
+                break;
+
+            case ParameterTranslationMode.Constant:
+                break;
+
+            default:
+                throw new UnreachableException();
+        }
+    }
+}

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

@@ -839,10 +839,42 @@ public override async Task Parameter_collection_Count_with_column_predicate_with
     public override Task Parameter_collection_Count_with_huge_number_of_values()
         => base.Parameter_collection_Count_with_huge_number_of_values();
 
+    // nothing to test here
+    public override Task Parameter_collection_Count_with_huge_number_of_values_over_5_operations()
+        => base.Parameter_collection_Count_with_huge_number_of_values_over_5_operations();
+
+    // nothing to test here
+    public override Task Parameter_collection_Count_with_huge_number_of_values_over_5_operations_forced_constants()
+        => base.Parameter_collection_Count_with_huge_number_of_values_over_5_operations_forced_constants();
+
+    // nothing to test here
+    public override Task Parameter_collection_Count_with_huge_number_of_values_over_5_operations_same_parameter()
+        => base.Parameter_collection_Count_with_huge_number_of_values_over_5_operations_same_parameter();
+
+    // nothing to test here
+    public override Task Parameter_collection_Count_with_huge_number_of_values_over_5_operations_mixed_parameters_constants()
+        => base.Parameter_collection_Count_with_huge_number_of_values_over_5_operations_mixed_parameters_constants();
+
     // nothing to test here
     public override Task Parameter_collection_of_ints_Contains_int_with_huge_number_of_values()
         => base.Parameter_collection_of_ints_Contains_int_with_huge_number_of_values();
 
+    // nothing to test here
+    public override Task Parameter_collection_of_ints_Contains_int_with_huge_number_of_values_over_5_operations()
+        => base.Parameter_collection_of_ints_Contains_int_with_huge_number_of_values_over_5_operations();
+
+    // nothing to test here
+    public override Task Parameter_collection_of_ints_Contains_int_with_huge_number_of_values_over_5_operations_same_parameter()
+        => base.Parameter_collection_of_ints_Contains_int_with_huge_number_of_values_over_5_operations_same_parameter();
+
+    // nothing to test here
+    public override Task Parameter_collection_of_ints_Contains_int_with_huge_number_of_values_over_5_operations_forced_constants()
+        => base.Parameter_collection_of_ints_Contains_int_with_huge_number_of_values_over_5_operations_forced_constants();
+
+    // nothing to test here
+    public override Task Parameter_collection_of_ints_Contains_int_with_huge_number_of_values_over_5_operations_mixed_parameters_constants()
+        => base.Parameter_collection_of_ints_Contains_int_with_huge_number_of_values_over_5_operations_mixed_parameters_constants();
+
     public override async Task Column_collection_of_ints_Contains()
     {
         await base.Column_collection_of_ints_Contains();