Simplifying LazyLinq

This is the fourth in a series of posts on LazyLinq, a wrapper to support lazy initialization and deferred disposal of a LINQ query context:

  1. Introducing LazyLinq: Overview
  2. Introducing LazyLinq: Internals
  3. Introducing LazyLinq: Queryability
  4. Simplifying LazyLinq
  5. Introducing LazyLinq: Lazy DataContext

As I was iterating on the proof of concept for LazyLinq, I always wanted to get rid of the TQuery type parameter. I thought I needed it to distinguish between ordered and unordered wrapped queries, but it just felt messy. The underlying provider mechanism didn’t need it, so why should I?

Well after taking a closer look at the SQL query provider, I figured out how to eliminate it. The object of inspiration is System.Data.Linq.DataQuery<T>, defined as follows:

internal sealed class DataQuery<T> :
    IOrderedQueryable<T>, IQueryable<T>,
    IOrderedQueryable, IQueryable,
    IEnumerable<T>, IEnumerable,
    IQueryProvider,
    IListSource

The key was realizing that IOrderedQueryable<> and ILazyOrderedQueryable<> don’t actually do anything. Implementation-wise, they’re just IQueryable<> or ILazyQueryable<> with an extra interface on top. It’s only on the design side that it actually matters, essentially providing a hook for additional ordering with ThenBy. In LINQ to SQL’s case, that means supporting orderability is as simple as specifying that the query object is both IQueryable<> and IOrderedQueryable<>.

So how does this revelation simplify Lazy LINQ? First, it allows us to remove TQuery from the interfaces:

    public interface ILazyContext<TContext> : IDisposable
    {
        TContext Context { get; }
        ILazyQueryable<TContext, TResult>
            CreateQuery<TResult>(Func<TContext, IQueryable<TResult>> queryBuilder);
        TResult Execute<TResult>(Func<TContext, TResult> action);
    }

    public interface ILazyQueryable<TContext, TSource> : IQueryable<TSource>
    {
        ILazyContext<TContext> Context { get; }
        Func<TContext, IQueryable<TSource>> QueryBuilder { get; }
    }

    public interface ILazyOrderedQueryable<TContext, TSource>
        : ILazyQueryable<TContext, TSource>, IOrderedQueryable<TSource>
    { }

Note that we can also eliminate ILazyContext.CreateOrderedQuery(), instead assuming that CreateQuery() will return something that can be treated as ILazyOrderedQueryable<> as necessary.

For the concrete implementations, we take the cue from DataQuery<T>, letting LazyQueryableImpl implement ILazyOrderedQueryable<> so we can eliminate LazyOrderedQueryableImpl:

    class LazyQueryableImpl<TContext, TSource>
        : ILazyQueryable<TContext, TSource>, ILazyOrderedQueryable<TContext, TSource>
    {
        // Implementation doesn't change
    }

Finally, our sorting query operations will look more like their counterparts in System.Linq.Queryable, casting the result of CreateQuery() to ILazyOrderedQueryable<>. To keep things readable, we’ll split our CreateOrderedQuery<> helper into separate versions for OrderBy and ThenBy. Note how the types of queryOperation map to the usage of OrderBy (unordered to ordered) and ThenBy (ordered to ordered):

        private static ILazyOrderedQueryable<TContext, TResult>
            CreateOrderByQuery<TSource, TContext, TResult>(
                this ILazyQueryable<TContext, TSource> source,
                Func<IQueryable<TSource>, IOrderedQueryable<TResult>> queryOperation
            )
        {
            return (ILazyOrderedQueryable<TContext, TResult>) source.Context.CreateQuery<TResult>(
               context => queryOperation(source.QueryBuilder(context)));
        }

        private static ILazyOrderedQueryable<TContext, TResult>
            CreateThenByQuery<TSource, TContext, TResult>(
                this ILazyQueryable<TContext, TSource> source,
                Func<IOrderedQueryable<TSource>, IOrderedQueryable<TResult>> queryOperation
            )
        {
            return (ILazyOrderedQueryable<TContext, TResult>) source.Context.CreateQuery<TResult>(
               context => queryOperation((IOrderedQueryable<TSource>) source.QueryBuilder(context)));
        }

Removing TQuery from the query operators is left as an exercise for the reader. Or you can just get the source on CodePlex.

A Note on IOrderedEnumerable<>

Having taken advantage of how LINQ to IQueryable handles orderability, it’s worth pointing out that LINQ to Objects uses a different approach, specifying new behavior in IOrderedEnumerable<> that is used to support multiple sort criteria:

public interface IOrderedEnumerable<TElement> : IEnumerable<TElement>, IEnumerable
{
    IOrderedEnumerable<TElement>
        CreateOrderedEnumerable<TKey>(
            Func<TElement, TKey> keySelector,
            IComparer<TKey> comparer, bool descending);
}

Introducing LazyLinq: Queryability

This is the third in a series of posts on LazyLinq, a wrapper to support lazy initialization and deferred disposal of a LINQ query context:

  1. Introducing LazyLinq: Overview
  2. Introducing LazyLinq: Internals
  3. Introducing LazyLinq: Queryability
  4. Simplifying LazyLinq
  5. Introducing LazyLinq: Lazy DataContext

Having defined the LazyLinq interfaces and provided concrete implementations, we’re left to provide support for the standard query operators.

Learning from Queryable

Before we try to query ILazyQueryable, it’s instructive to look at how System.Linq.Queryable works. There are essentially three types of operators on IQueryable<>:

  • Deferred queries returning IQueryable<>: Select, Where, etc.
  • Deferred query returning IOrderedQueryable<>: OrderBy, OrderByDescending, ThenBy, ThenByDescending
  • Everything else: Aggregate, Count, First, etc.

Reflecting Queryable.Select, modulo error checking, we see the following:

public static IQueryable<TResult> Select<TSource, TResult>(
    this IQueryable<TSource> source, Expression<Func<TSource, TResult>> selector)
{
    return source.Provider
        .CreateQuery<TResult>(
            Expression.Call(null,
                ((MethodInfo) MethodBase.GetCurrentMethod())
                    .MakeGenericMethod(new Type[] { typeof(TSource), typeof(TResult) }),
                new Expression[] { source.Expression, Expression.Quote(selector) }));
}

The source‘s Provider is used to construct a new query whose expression includes the call to Select with the given parameters. An ordered query follows a similar pattern, trusting that the query provider will return an IOrderedQueryable<> as appropriate:

public static IOrderedQueryable<TSource> OrderBy<TSource, TKey>(
    this IQueryable<TSource> source, Expression<Func<TSource, TKey>> keySelector)
{
    return (IOrderedQueryable<TSource>) source.Provider
        .CreateQuery<TSource>(
            Expression.Call(null,
                ((MethodInfo) MethodBase.GetCurrentMethod())
                    .MakeGenericMethod(new Type[] { typeof(TSource), typeof(TKey) }),
                new Expression[] { source.Expression, Expression.Quote(keySelector) }));
}

And finally, everything that’s not a query is handled by the provider’s Execute method:

public static TSource First<TSource>(this IQueryable<TSource> source)
{
    return source.Provider
        .Execute<TSource>(
            Expression.Call(null,
                ((MethodInfo) MethodBase.GetCurrentMethod())
                    .MakeGenericMethod(new Type[] { typeof(TSource) }),
                new Expression[] { source.Expression }));
}

Querying ILazyQueryable

You may have noticed that the above scenarios map rather closely to the methods provided by ILazyContext:

  • ILazyQueryable<TContext, TResult, TQuery>
        CreateQuery<TResult, TQuery>(Func<TContext, TQuery> queryBuilder)
        where TQuery : IQueryable<TResult>;
  • ILazyOrderedQueryable<TContext, TResult, TQuery>
        CreateOrderedQuery<TResult, TQuery>(Func<TContext, TQuery> queryBuilder)
        where TQuery : IOrderedQueryable<TResult>;
  • TResult Execute<TResult>(Func<TContext, TResult> action);

However, rather than expression trees we’re pushing around delegates. Execute seems pretty simple, so let’s start there:

public static TSource First<TSource, TContext, TQuery>(
        this ILazyQueryable<TContext, TSource, TQuery> source
    ) where TQuery : IQueryable<TSource>
{
    Func<TContext, TResult> action = context => ???;
    return source.Context.Execute(action);
}

So our source has a Context, which knows how to Execute an action from context to some result. To find that result, we need to leverage the other property of source: QueryBuilder. Recalling that QueryBuilder is a function from TContext to TQuery, and that TQuery is an IQueryable<TSource>, we see something on which we can execute Queryable.First:

public static TSource First<TSource, TContext, TQuery>(
        this ILazyQueryable<TContext, TSource, TQuery> source
    ) where TQuery : IQueryable<TSource>
{
    Func<TContext, TResult> action =
        context => source.QueryBuilder(context).First();
    return source.Context.Execute(action);
}

Now seeing as we have dozens of methods to implement like this, it seems an opportune time for a bit of eager refactoring. Recognizing that the only variance is the method call on the IQueryable, let’s extract an extension method that does everything else:

private static TResult Execute<TSource, TContext, TResult, TQuery>(
        this ILazyQueryable<TContext, TSource, TQuery> source,
        Func<TQuery, TResult> queryOperation
    ) where TQuery : IQueryable<TSource>
{
    return source.Context.Execute(
        context => queryOperation(source.QueryBuilder(context)));
}

From there, additional lazy operators are just a lambda expression away:

public static TSource First<TSource, TContext, TQuery>(
        this ILazyQueryable<TContext, TSource, TQuery> source
    ) where TQuery : IQueryable<TSource>
{
    return source.Execute(q => q.First());
}

public static TAccumulate Aggregate<TContext, TSource, TQuery, TAccumulate>(
        this ILazyQueryable<TContext, TSource, TQuery> source,
        TAccumulate seed, Expression<Func<TAccumulate, TSource, TAccumulate>> func
    ) where TQuery : IQueryable<TSource>
{
    return source.Execute(q => q.Aggregate(seed, func));
}

And now having done the hard part (finding an IQueryable), we can translate that understanding to make similar helpers for queries:

private static ILazyQueryable<TContext, TResult, IQueryable<TResult>>
    CreateQuery<TSource, TContext, TQuery, TResult>(
        this ILazyQueryable<TContext, TSource, TQuery> source,
        Func<TQuery, IQueryable<TResult>> queryOperation
    ) where TQuery : IQueryable<TSource>
{
    return source.Context.CreateQuery<TResult, IQueryable<TResult>>(
        context => queryOperation(source.QueryBuilder(context)));
}

private static ILazyOrderedQueryable<TContext, TResult, IOrderedQueryable<TResult>>
    CreateOrderedQuery<TSource, TContext, TQuery, TResult>(
        this ILazyQueryable<TContext, TSource, TQuery> source,
        Func<TQuery, IOrderedQueryable<TResult>> queryOperation
    ) where TQuery : IQueryable<TSource>
{
    return source.Context.CreateOrderedQuery<TResult, IOrderedQueryable<TResult>>(
        context => queryOperation(source.QueryBuilder(context)));
}

With similarly trivial query operator implementations:

public static ILazyQueryable<TContext, TResult, IQueryable<TResult>>
    Select<TContext, TSource, TQuery, TResult>(
        this ILazyQueryable<TContext, TSource, TQuery> source,
        Expression<Func<TSource, TResult>> selector
    ) where TQuery : IQueryable<TSource>
{
    return source.CreateQuery(q => q.Select(selector));
}

public static ILazyOrderedQueryable<TContext, TSource, IOrderedQueryable<TSource>>
    OrderBy<TContext, TSource, TQuery, TKey>(
        this ILazyQueryable<TContext, TSource, TQuery> source,
        Expression<Func<TSource, TKey>> keySelector
    ) where TQuery : IQueryable<TSource>
{
    return source.CreateOrderedQuery(q => q.OrderBy(keySelector));
}

And the end result:

LazyLinq.First

Introducing LazyLinq: Internals

This is the second in a series of posts on LazyLinq, a wrapper to support lazy initialization and deferred disposal of a LINQ query context:

  1. Introducing LazyLinq: Overview
  2. Introducing LazyLinq: Internals
  3. Introducing LazyLinq: Queryability
  4. Simplifying LazyLinq
  5. Introducing LazyLinq: Lazy DataContext

My first post introduced the three interfaces that LazyLinq provides. Next, we get to implement them.

Implementing ILazyQueryable

First, the interface:

    public interface ILazyQueryable<TContext, TSource, TQuery>
        : IQueryable<TSource>
        where TQuery : IQueryable<TSource>
    {
        ILazyContext<TContext> Context { get; }
        Func<TContext, TQuery> QueryBuilder { get; }
    }

We’ll start simple with an implicit implementation of the interface and a trivial constructor:

    class LazyQueryableImpl<TContext, TSource, TQuery>
            : ILazyQueryable<TContext, TSource, TQuery>
            where TQuery : IQueryable<TSource>
    {
            public ILazyContext<TContext> Context { get; private set; }
            public Func<TContext, TQuery> QueryBuilder { get; private set; }

            internal LazyQueryableImpl(ILazyContext<TContext> deferredContext, Func<TContext, TQuery> queryBuilder)
            {
                if (deferredContext == null) throw new ArgumentNullException("deferredContext");
                if (queryBuilder == null) throw new ArgumentNullException("queryBuilder");

                Context = deferredContext;
                QueryBuilder = queryBuilder;
            }

Next, a lazy-loaded query built from our lazy context:

            protected TQuery Query
            {
                get
                {
                    if (query == null)
                    {
                        query = QueryBuilder(Context.Context);
                        if (query == null)
                            throw new InvalidOperationException("Query built as null.");
                    }
                    return query;
                }
            }

And the internals of managing Context, which implements IDisposable:

            private void Dispose()
            {
                Context.Dispose();
                query = default(TQuery);
            }

            private IEnumerator<TSource> GetEnumerator()
            {
                try
                {
                    foreach (var i in Query)
                        yield return i;
                }
                finally
                {
                    Dispose();
                }
            }

Since Query depends on Context, once Context is disposed we need to reset Query so a new one can be built (if possible). Note that we use an iterator here to return an IEnumerator<TSource>, rather than the usual IEnumerable<>.

Finally, we’ll close out by explicitly implementing IQueryable:

            IEnumerator<TSource> IEnumerable<TSource>.GetEnumerator()
            {
                return GetEnumerator();
            }

            IEnumerator IEnumerable.GetEnumerator()
            {
                return GetEnumerator();
            }

            Type IQueryable.ElementType
            {
                get { return Query.ElementType; }
            }

            Expression IQueryable.Expression
            {
                get { return Query.Expression; }
            }

            IQueryProvider IQueryable.Provider
            {
                get { return Query.Provider; }
            }
        }

If this seemed relatively simple, you’re right. We’re just building a lazy-loaded Query proxy, with a bit of plumbing to clean up our Context.

Implementing ILazyOrderedQueryable

Not very exciting, but for completeness:

        public interface ILazyOrderedQueryable<TContext, TSource, TQuery>
            : ILazyQueryable<TContext, TSource, TQuery>, IOrderedQueryable<TSource>
            where TQuery : IOrderedQueryable<TSource>
        { }

        class LazyOrderedQueryableImpl<TContext, TSource, TQuery>
                : LazyQueryableImpl<TContext, TSource, TQuery>, ILazyOrderedQueryable<TContext, TSource, TQuery>
                where TQuery : IOrderedQueryable<TSource>
        {
            internal LazyOrderedQueryableImpl(ILazyContext<TContext> lazyContext, Func<TContext, TQuery> queryBuilder)
                : base(lazyContext, queryBuilder)
            {
            }
        }

LazyQueryable Factory

Consumers of this API should never need to know about these implementation details, so we can hide them behind a factory class:

        public static class LazyQueryable
        {
            public static ILazyQueryable<TContext, TResult, TQuery> CreateQuery<TContext, TResult, TQuery>(
                ILazyContext<TContext> context, Func<TContext, TQuery> queryBuilder)
                where TQuery : IQueryable<TResult>
            {
                return new LazyQueryableImpl<TContext, TResult, TQuery>(context, queryBuilder);
            }
            public static ILazyOrderedQueryable<TContext, TResult, TQuery> CreateOrderedQuery<TContext, TResult, TQuery>(
                ILazyContext<TContext> context, Func<TContext, TQuery> queryBuilder)
                where TQuery : IOrderedQueryable<TResult>
            {
                return new LazyOrderedQueryableImpl<TContext, TResult, TQuery>(context, queryBuilder);
            }
        }

Implementing ILazyContext

Again, we’ll start with the interface:

        public interface ILazyContext<TContext> : IDisposable
        {
            TContext Context { get; }

            ILazyQueryable<TContext, TResult, TQuery>
                CreateQuery<TResult, TQuery>(Func<TContext, TQuery> queryBuilder)
                where TQuery : IQueryable<TResult>;

            ILazyOrderedQueryable<TContext, TResult, TQuery>
                CreateOrderedQuery<TResult, TQuery>(Func<TContext, TQuery> queryBuilder)
                where TQuery : IOrderedQueryable<TResult>;

            TResult Execute<TResult>(Func<TContext, TResult> action);
        }

Now we can start fulfilling our requirements:

1. Lazily expose the Context.

        class LazyContextImpl<TContext> : ILazyContext<TContext>, IDisposable
        {
            public Func<TContext> ContextBuilder { get; private set; }

            private TContext context;
            public TContext Context
            {
                get
                {
                    if (context == null)
                    {
                        context = ContextBuilder();
                        if (context == null)
                            throw new InvalidOperationException("Context built as null.");
                    }
                    return context;
                }
            }

2. Produce lazy wrappers to represent queries retrieved from a context by a delegate.

            public ILazyQueryable<TContext, TResult, TQuery> CreateQuery<TResult, TQuery>(
                Func<TContext, TQuery> queryBuilder)
                where TQuery : IQueryable<TResult>
            {
                return LazyQueryable.CreateQuery<TContext, TResult, TQuery>(this, queryBuilder);
            }

            public ILazyOrderedQueryable<TContext, TResult, TQuery> CreateOrderedQuery<TResult, TQuery>(
                Func<TContext, TQuery> queryBuilder)
                where TQuery : IOrderedQueryable<TResult>
            {
                return LazyQueryable.CreateOrderedQuery<TContext, TResult, TQuery>(this, queryBuilder);
            }

3. Execute an action on the context.

There are two ways to “complete” a query, and we need to clean up context after each. The first was after enumeration, implemented above. The second is on execute, implemented here:

            public TResult Execute<TResult>(Func<TContext, TResult> expression)
            {
                try
                {
                    return expression(Context);
                }
                finally
                {
                    Dispose();
                }
            }

4. Ensure the context is disposed as necessary.

We don’t require that TContext is IDisposable, but we need to handle if it is. We also clear context to support reuse.

            public void Dispose()
            {
                var disposable = context as IDisposable;
                if (disposable != null)
                    disposable.Dispose();

                context = default(TContext);
            }

Constructors

With our requirements met, we just need a way to create our context. We provide two options:

            internal LazyContextImpl(TContext context) : this(() => context) { }

            internal LazyContextImpl(Func<TContext> contextBuilder)
            {
                if (contextBuilder == null) throw new ArgumentNullException("contextBuilder");

                ContextBuilder = contextBuilder;
            }

The former wraps an existing TContext instance in a closure, meaning every time ContextBuilder is called it returns the same instance. The latter accepts any delegate that returns a TContext. The most common such delegate would be a simple instantiation: () => new MyDataContext().

It should be clear now why we would want to clear our context on dispose. If ContextBuilder returns a new context instance each time, it’s perfectly safe to discard of the old (disposed) context to trigger the creation of a new one. Conversely, if the builder returns a single instance, using the context after disposal would trigger an ObjectDisposedException or something similar.

LazyContext Factory

For consistency, we should also provide factory methods to hide this specific implementation:

            public static class LazyContext
            {
                public static ILazyContext<T> Create<T>(T context)
                {
                    return new LazyContextImpl<T>(context);
                }
                public static ILazyContext<T> Create<T>(Func<T> contextBuilder)
                {
                    return new LazyContextImpl<T>(contextBuilder);
                }
            }

Lazy Extensions

And last, but certainly not least, we’re ready to reimplement our Use() extension methods:

    public static class Lazy
    {
        public static ILazyContext<TContext> Use<TContext>(this TContext @this)
        {
            return LazyContext.Create<TContext>(@this);
        }

        public static ILazyContext<TContext> Use<TContext>(this Func<TContext> @this)
        {
            return LazyContext.Create<TContext>(@this);
        }
    }

With several usage possibilities:

    var r1 = from x in new MyDataContext().Use() ...;

    Func<MyDataContext> f1 = () => new MyDataContext();
    var r2 = from x in f1.Use() ...;

    var r3 = from x in new Func<MyDataContext>(() => new MyDataContext()).Use() ...;

    var r4 = from x in Lazy.Use(() => new MyDataContext()) ...;

Or maybe we can make it even easier. Maybe…