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;
        }
    }
}

I read that answer, but did not understand why a bucket size of 32 performs better than one of size 20.
A sql plan is always "created" in sizes of multiple of 2?

@juliowh It is the number of unique query plans that matter, and a factor 2 seemed like suitable bucket sizes

That's very cool @ErikEJ !

We use a lot of composite keys and I wonder if would be possible to create another overload that works with composite keys?

So I guess the signature might be something like this:

public static IQueryable<TQuery> In<TKey, TQuery>(
this IQueryable<TQuery> queryable,
IEnumerable<Tuple<TKey1, TKey2>> values,
Expression<Func<TQuery, Tuple<TKey1, TKey2>>> keySelector)

And we might use it like this:

[PrimaryKey(nameof(State), nameof(LicensePlate))]
internal class Car
{
    public string State { get; set; }
    public string LicensePlate { get; set; }

    public string Make { get; set; }
    public string Model { get; set; }
}

var keys = new[] { ("state1", "license1"), ("state2", "license2"), etc... }
var cars = context.Cars.In(keys, c => (c.State, c.LicensePlate));

It's a bit more complicated because the predicate needs to operate on two separate columns to generate something like this:

SELECT ...
FROM ...
WHERE (State = @pState1 AND LicensePlate = @pLicensePlate1)
OR    (State = @pState2 AND LicensePlate = @pLicensePlate2)
OR ...

Related: dotnet/efcore#32092

@clement911 Feel free to do with this snippet whatever you want. It is just a sample.