perf: index-aware LINQ terminals with predicates, single-copy page writes, single-pass B-Tree insert (#179)

[CaffeinatedCoder]

Fixes #179.

Four small changes from profiling BLite against the SQLite benchmarks: two query-side, two on the write path.

Changes

1 + 2. LINQ terminals with a predicate use the index

• BTreeExpressionVisitor folds the predicate of Count/LongCount and of First, FirstOrDefault, Single, SingleOrDefault and Any into the WHERE clause, using the same AndAlso merge as chained Where calls.
• This happens only when every operator between the collection and the terminal is a single-parameter Where, OrderBy* or ThenBy*, and the predicate takes the entity type. After anything else (Select, Take, Skip, SkipWhile, TakeWhile, GroupBy, OfType, Cast, SelectMany, Join, Zip, the indexed Where) the query stays on the fallback pipeline. Filtering commutes with ordering, so an OrderBy before it is fine. All is never folded. The .NET 6+ overloads that take a default value are not folded either.
• When only a WHERE is left in front of an element terminal, the provider fetches one document (two for Single*) through the index-aware FetchAsync and applies the terminal to that list. Otherwise the query is marked complex exactly as before; EnumerableRewriter re-applying an already applied predicate is harmless.

3. One copy per page write. Transaction.AddWrite passes NewValue.Span straight to StorageEngine.WritePage instead of ToArray()-ing it first. WritePage already copies into the transaction's own page buffer (data.ToArray() for a new page, CopyTo for an existing one), so callers can still return pooled buffers right away. A new test covers this.

4. Single-pass unique B-Tree insert

• BTreeIndex.Insert finds the leaf once (FindLeafNodeWithPath now has an overload that leaves the leaf in the caller's buffer) and does the duplicate check with BinarySearchLeaf on that leaf. Before, a separate TryFind walked from the root first, and the leaf was read a third time.
• Duplicate semantics are unchanged: the same key with the same location is a no-op, the same key with a different location throws.
• TryFind rents a pooled buffer instead of zeroing a 16 KB stackalloc on every lookup.
• The leaf shift uses Span.CopyTo (memmove) instead of a byte loop.

Behaviour changes

• Unordered First(pred), FirstOrDefault(pred) and Any(pred) follow index order. When an index serves the predicate, the result is the first match in index order, not in collection scan order. This is what Where(pred).First() already did, and LINQ promises no order without an OrderBy. Code that relied on scan order should add an OrderBy.
• Count(p => p.V.VectorSearch(q, k)) returns k (at most the collection size), the same as Where(p => p.V.VectorSearch(q, k)).Count(). Before, the predicate was evaluated per document in memory, and the count covered the whole collection.

Measurements

macOS arm64, Release, averages of a warm loop over 6,000 PersonV2 documents with idx_Age, after a forced GC:

	before	after
Count(p => p.Age == x)	1,663 µs	15 µs
Count(p => p.Age >= a && p.Age <= b)	1,749 µs	73 µs
FirstOrDefault(p => p.Age == x)	1,917 µs	9 µs
Any(p => p.Age == x)	1,464 µs	14 µs
FirstOrDefault(p => p.Name == x) (no index, BSON scan stops early)	1,418 µs	724 µs
SingleOrDefault(p => p.Id == id)	2,174 µs	860 µs
Count(p => p.Name == x) (no index)	2,599 µs	1,243 µs
1,000 unique B-Tree inserts, one transaction	3.5 ms	1.6 ms
InsertBulkAsync 1,000 docs (pk + idx_Age)	14.2 ms / 22.3 MB	10.7 ms / 6.2 MB

FindByIdAsync is within noise in this harness, about 3.4–3.9 µs either way. A tighter loop earlier measured 4.0 → 3.15 µs from the TryFind change.

Tests

QueryAndWritePathTests adds 13 tests:

• Query model: which forms get folded, and that the predicate is not folded after Take, Skip, Select, SkipWhile or the indexed Where.
• Terminals with a predicate after GroupBy, SkipWhile, TakeWhile, OfType and Cast match LINQ-to-Objects (these crashed or returned wrong results with the first version of this PR).
• Count(x => x.Embedding.VectorSearch(q, 3)) returns 3.
• Count/LongCount with indexed, range and non-indexed predicates, and OrderBy().Take().Count(p), all compared against LINQ-to-Objects.
• Every element terminal with and without matches, including the exceptions for First with no match and Single/SingleOrDefault with several.
• Terminals after OrderBy, Where, Skip and Select, plus All, and the default-value overload of FirstOrDefault.
• A caller's buffer can be reused immediately after AddWrite.
• The single-pass unique insert rejects duplicates and ignores exact re-inserts.

Full suite on macOS: 2392 passed, 3 failed, 4 skipped. The 3 failures are the known MultiProcessWalSharedMemoryTests .wal-shm lock timeouts. The branch merges cleanly with #169, #170, #172, #173, #174, #175 and #177.

Found along the way, not changed here

After a Select, scalar terminals return the default value. On main, Select(p => p.Age).Any(), .Any(pred), .Count(pred), .Max() and .Sum() return false/0 on a non-empty collection. ProjectEnumerable returns default when the result type isn't a sequence. That is a correctness bug, so I'll report it separately.

🤖 Generated with Claude Code

[mrdevrobot]

Thanks. The single-copy page write is safe: StorageEngine.WritePage always copies, and every AddWrite caller passes a pooled buffer only for the duration of the call. The single-pass B-Tree insert behaves exactly like before, since TryFind and the insert path use the same leaf routing. The terminal semantics also check out against LINQ-to-Objects. The problem is when the predicate is folded.

1. The folded predicate crashes after GroupBy (BTreeExpressionVisitor.cs:171-176, :197-202). The guard checks SelectClause, Take and Skip, but not HasComplexOperators or the element type. Q.GroupBy(p => p.Age).Count(g => g.Count() > 5) and .Any(g => g.Key == 3) put a lambda over IGrouping<,> into WhereClause, and FetchAsync throws InvalidCastException when it casts it to Expression<Func<T,bool>>. On main these return 50 and true. OfType<>, Cast<>, SelectMany, Join and Zip are exposed the same way.

2. Wrong results after SkipWhile / TakeWhile. Q.SkipWhile(p => p.Age < 10).FirstOrDefault(p => p.Age == 5) returns null; main and LINQ-to-Objects return the matching document. The predicate is pushed below an order-dependent operator.

Fix for both: fold only when !_model.HasComplexOperators and predicate.Parameters[0].Type == typeof(T), or allowlist the operators that may come before the terminal (Where, OrderBy*, ThenBy*). Please add regression tests for GroupBy, SkipWhile and OfType.

3. Behaviour changes to mention in the description. Unordered First(pred) now returns the first match in index order, not scan order. That is legal and matches Where(p).First(), but users can see it. Count(p => p.V.VectorSearch(q, k)) now returns k instead of the whole collection; that's better, but untested.

This PR doesn't touch #181. If you're in this code anyway, skipping TryPushDownSelect for element terminals when the push-down doesn't apply would fix part of it cheaply; a separate PR is fine too.

Nits: the XML doc at Transaction.cs:51-55 still says "Makes a defensive copy"; DocumentCollection.cs:1854 still does a redundant ToArray() before AddWrite. Merges cleanly with main, #169 and #173.

[CaffeinatedCoder]

Thanks. Addressed in 1f4dd2f.

1 + 2 (predicate folded after GroupBy / SkipWhile / …): I went with the allowlist. CanFoldPredicate walks the source chain and folds the terminal's predicate only when every operator between the collection root and the terminal is a single-parameter Where, OrderBy* or ThenBy*, and the predicate's parameter type is the root's element type. Anything else (GroupBy, OfType, Cast, Select, SelectMany, Join, Zip, Take, Skip, SkipWhile, TakeWhile, the indexed Where((p, i) => ...)) leaves the predicate on the fallback path. This replaces the old Select/Take/Skip guard and covers it.

Regression tests, each failing on 73e9c99:

• Terminals_AfterGroupBy_PredicateIsOverGroups: GroupBy(...).Count(g => ...), .Any(g => g.Key == 3), .First(g => ...)
• Terminals_AfterSkipWhile_PredicateIsNotPushedBelowIt and Terminals_AfterTakeWhile_..., compared against the collection's own scan order, since that is what these operators depend on
• Terminals_AfterOfTypeOrCast_PredicateIsOverTheNewElementType: OfType<object>().Count(...), Cast<object>().Any(...)
• Model_PredicateIsFoldedOnlyAfterWhereAndOrdering: a model-level check of the allowlist

3: the description now has a "Behaviour changes" section with both points. Count(p => p.V.VectorSearch(q, k)) returning k is covered by Count_WithVectorSearchPredicate_CountsTheNearestK.

#181: left for a separate PR.

Nits: the Transaction.AddWrite XML doc is updated, and the ToArray() in AllocateOverflowPage is gone.

Related, pre-existing on main: a plain Where after these operators has the same push-down problem, because VisitWhere always merges into WhereClause. GroupBy(...).Where(g => ...) and Where((p, i) => ...) throw InvalidCastException, and SkipWhile(...).Where(...) / Take(n).Where(...) return wrong results. It is the same root cause as the Join inner Where from #173, so both are in #186. The CanFoldPredicate check could be reused in VisitWhere.

Full suite: 2392 passed, 3 failed (the known MultiProcessWalSharedMemoryTests .wal-shm timeouts, which also fail without this commit), 4 skipped.

[mrdevrobot]

Verified 1f4dd2f. I compared 28 query shapes against LINQ-to-Objects over the same 300 rows: GroupBy, SkipWhile, TakeWhile, OfType, Cast, Select, Where + OrderBy/ThenBy before First, Take/Skip/Distinct/Reverse before a predicate terminal, Single/SingleOrDefault/Last. All match, and ordering is respected when an OrderBy precedes the terminal. The only mismatch is the indexed Where((p, i) => ...), which throws on main too (#186). Five of the new tests fail with the fix reverted.

One thing before merging. #173 is now on main, and combined with this PR one of its tests fails:

NearWithinCorrectnessTests.VectorSearch_InCountPredicate_AfterOtherWhere_Throws
Assert.Throws() Failure: No exception was thrown

Where(x => x.Title != null).Count(x => x.Embedding.VectorSearch(Q, 1)) now has its predicate folded into the WHERE, so it runs as Where(Title != null && VectorSearch(Q, 1)).Count(): the vector term drives and Title != null is the post-filter. That's consistent with how #173 treats VectorSearch && x, and with your Count_WithVectorSearchPredicate_CountsTheNearestK, so I think the new behaviour is the right one and the test should change, not the code. Please rebase onto main and update that test to assert the folded result (the same count as Where(x => x.Title != null && x.Embedding.VectorSearch(Q, 1)).Count()). The Sum(x => VectorSearch(...) ? 1 : 0) variant isn't folded and still throws. With that, the combination of main and this branch is otherwise green.

[CaffeinatedCoder]

Rebased onto main (efacaf4). Both commits replayed without conflicts and are unchanged (73e9c99 → ae6308e, 1f4dd2f → ef05b29).

• a167c72: renamed VectorSearch_InCountPredicate_AfterOtherWhere_Throws to ..._IsFoldedIntoTheWhere. It now asserts that Where(x => x.Title != null).Count(x => x.Embedding.VectorSearch(Q, 1)) returns the same count as Where(x => x.Title != null && x.Embedding.VectorSearch(Q, 1)).Count() (1 on the test data). It fails with the folding reverted.
• VectorSearch_InSumSelector_AfterOtherWhere_Throws is unchanged and still passes, so the Sum(x => VectorSearch(...) ? 1 : 0) variant still throws.
• Full BLite.Tests: 2470 / 0 / 4 skipped.

One related shape: OrderBy(x => x.Title).First(x => x.Embedding.VectorSearch(Q, 1)) still throws. The predicate is folded into the WHERE, but because of the OrderBy the element terminal goes through the fallback, which replays First(p) and evaluates the VectorSearch again. It throws on main as well, so it isn't a regression, but it doesn't match OrderBy(...).Count(VectorSearch), which returns the folded result. If you'd like them to match, I can drop the folded terminal predicate from the replay (the same idea as DropAppliedWheres), here or in a follow-up.

[CaffeinatedCoder]

Follow-up in 3cba537, which also covers the OrderBy(...).First(VectorSearch) shape from my last comment: a terminal predicate that was folded into the WHERE is no longer evaluated again in the fallback replay.

• OrderBy(x => x.Title).First(x => x.Embedding.VectorSearch(Q, 1)) now returns the same result as Where(x => x.Embedding.VectorSearch(Q, 1)).OrderBy(x => x.Title).First(). The replay swaps the folded terminal for its parameterless overload (First(p) becomes First()), next to DropAppliedWheres. This only happens when folding was allowed, i.e. only Where/OrderBy/ThenBy in front, so every element reaching the terminal has already passed p.
• This also fixes FirstOrDefault/Single/SingleOrDefault/Any after OrderBy, and Count/LongCount after ThenBy. ThenBy isn't modelled by the direct pipeline, so those queries reached the replay too.
• Last/LastOrDefault now fold their predicate as well (Last(p) is Where(p).Last()). They always go through the fallback, because they need every match.
• All(p) is not folded, since it isn't Where(p).All(), and it is still evaluated in full, so All(VectorSearch) still throws. Predicates after SkipWhile/Take/GroupBy, and the Sum selector, are still evaluated in the replay and still throw.
• New tests in QueryAndWritePathTests cover OrderBy/ThenBy with each of these terminals using VectorSearch, compared with the equivalent Where(...) query; the same terminals with plain predicates against LINQ-to-Objects (including All true/false); and the unfolded shapes that must still throw. With the change reverted, every vector shape throws.
• Full BLite.Tests: 2474 / 0 / 4 skipped.

[mrdevrobot]

Verified the rebase and the two new commits. ae6308e and ef05b29 are identical to the commits I reviewed. a167c72 asserts the folded count against Where(Title != null && VectorSearch(Q, 1)).Count(), and the other #173 throw tests still pass. 3cba537 only swaps a terminal for its parameterless overload when CanFoldPredicate allowed the fold, and the fetch then applies the whole predicate, partial-AND residuals included. Probes against LINQ-to-Objects match, including OrderBy(Title).First(VectorSearch(Q, 1)), and Select(x => x).First(VectorSearch) still throws. Merged with today's main, the full suite passes (2507 / 0 / 4 skipped). Merging.