Replacement for EF Core .Contains, that avoids SQL Server plan cache pollution

IQueryableExtensions.cs

using System.Linq.Expressions;

namespace Microsoft.EntityFrameworkCore
{
    public static class IQueryableExtensions
    {
        public static IQueryable<TQuery> In<TKey, TQuery>(
        this IQueryable<TQuery> queryable,
        IEnumerable<TKey> values,
        Expression<Func<TQuery, TKey>> keySelector)
        {
            ArgumentNullException.ThrowIfNull(values);

            ArgumentNullException.ThrowIfNull(keySelector);

            if (!values.Any())
            {
                return queryable.Take(0);
            }

            var distinctValues = Bucketize(values);

            if (distinctValues.Length > 2048)
            {
                throw new ArgumentException("Too many parameters for SQL Server, reduce the number of parameters", nameof(keySelector));
            }

            var expr = CreateBalancedORExpression(distinctValues, keySelector.Body, 0, distinctValues.Length - 1);

            var clause = Expression.Lambda<Func<TQuery, bool>>(expr, keySelector.Parameters);

            return queryable.Where(clause);
        }

        private static BinaryExpression CreateBalancedORExpression<TKey>(TKey[] values, Expression keySelectorBody, int start, int end)
        {
            if (start == end)
            {
                var v1 = values[start];
                return Expression.Equal(keySelectorBody, ((Expression<Func<TKey>>)(() => v1)).Body);
            }
            else if (start + 1 == end)
            {
                var v1 = values[start];
                var v2 = values[end];

                return Expression.OrElse(
                    Expression.Equal(keySelectorBody, ((Expression<Func<TKey>>)(() => v1)).Body),
                    Expression.Equal(keySelectorBody, ((Expression<Func<TKey>>)(() => v2)).Body));
            }
            else
            {
                int mid = (start + end) / 2;
                return Expression.OrElse(
                    CreateBalancedORExpression(values, keySelectorBody, start, mid),
                    CreateBalancedORExpression(values, keySelectorBody, mid + 1, end));
            }
        }

        private static TKey[] Bucketize<TKey>(IEnumerable<TKey> values)
        {
            var distinctValues = new HashSet<TKey>(values).ToArray();
            var originalLength = distinctValues.Length;

            int bucket = (int)Math.Pow(2, Math.Ceiling(Math.Log(originalLength, 2)));

            if (originalLength == bucket) return distinctValues;

            var lastValue = distinctValues[originalLength - 1];
            Array.Resize(ref distinctValues, bucket);
            distinctValues.AsSpan().Slice(originalLength).Fill(lastValue);

            return distinctValues;
        }
    }
}

@julienFlexsoft Thanks. It was updated in 2023, so I think it will work with current EF Core.

I simplified the CreateBalancedORExpression like this

private static BinaryExpression CreateBalancedORExpression<TKey>(TKey[] values, Expression keySelectorBody)
{
    return values
        .Select(p => Expression.Equal(keySelectorBody, ((Expression<Func<TKey>>)(() => p)).Body))
        .Aggregate(Expression.OrElse);
}

The resulting SQL looks almost identical, but it makes the parameter names a bit clearer. No real added value except the shorter code.

-- Old SQL
[Parameters=[@__v1_0='1', @__v2_1='2', @__v1_2='3', @__v2_3='3']]
SELECT [a].[Id]
FROM [ActivityEntities] AS [a]
WHERE [a].[Id] = @__v1_0 OR [a].[Id] = @__v2_1 OR [a].[Id] = @__v1_2 OR [a].[Id] = @__v2_3

-- New SQL
[Parameters=[@__p_0='1', @__p_1='2', @__p_2='3', @__p_3='3']]
SELECT [a].[Id]
FROM [ActivityEntities] AS [a]
WHERE [a].[Id] = @__p_0 OR [a].[Id] = @__p_1 OR [a].[Id] = @__p_2 OR [a].[Id] = @__p_3

What do you think?