From ae6308ec3770dc171d6968ec8d1b4b1fe02c0978 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ricardo=20Gro=C3=9F?= Date: Mon, 12 Oct 2026 15:07:27 +0200 Subject: [PATCH 1/4] perf: index-aware LINQ terminals with predicates, single-copy page writes, single-pass B-Tree insert (#179) - Count/LongCount and First/FirstOrDefault/Single/SingleOrDefault/Any with a predicate fold it into the WHERE clause when no Select/Take/Skip comes before, so the index planner sees it. With only a WHERE in front, element terminals fetch one document (two for Single*) instead of loading every match and running EnumerableRewriter. - Transaction.AddWrite passes the page straight to StorageEngine.WritePage, which already copies it, instead of allocating a 16 KB defensive copy per write. - Unique B-Tree inserts check for duplicates on the leaf they already found instead of a separate TryFind traversal, and reuse that leaf buffer. TryFind uses a pooled buffer instead of a zeroed 16 KB stackalloc, and leaf inserts shift entries with Span.CopyTo. Co-Authored-By: Claude Opus 5.5 --- src/BLite.Core/Indexing/BTreeIndex.cs | 87 +++++---- .../Query/BTreeExpressionVisitor.cs | 47 ++++- src/BLite.Core/Query/BTreeQueryProvider.cs | 27 +++ src/BLite.Core/Query/QueryModel.cs | 7 + src/BLite.Core/Transactions/Transaction.cs | 7 +- tests/BLite.Tests/QueryAndWritePathTests.cs | 167 ++++++++++++++++++ 6 files changed, 297 insertions(+), 45 deletions(-) create mode 100644 tests/BLite.Tests/QueryAndWritePathTests.cs diff --git a/src/BLite.Core/Indexing/BTreeIndex.cs b/src/BLite.Core/Indexing/BTreeIndex.cs index db20278c..d7ac1e29 100644 --- a/src/BLite.Core/Indexing/BTreeIndex.cs +++ b/src/BLite.Core/Indexing/BTreeIndex.cs @@ -101,26 +101,27 @@ public void Insert(IndexKey key, DocumentLocation location, ulong? transactionId { var txnId = transactionId ?? 0; - if (_options.Unique && TryFind(key, out var existingLocation, txnId)) - { - if (!existingLocation.Equals(location)) - throw new InvalidOperationException("Duplicate key violation for unique index"); - - return; - } - var entry = new IndexEntry(key, location); var path = new List(); - // Find the leaf node for insertion - var leafPageId = FindLeafNodeWithPath(key, path, txnId); - var pageBuffer = System.Buffers.ArrayPool.Shared.Rent(_storage.PageSize); try { - ReadPage(leafPageId, txnId, pageBuffer); + // Find the leaf node for insertion; the leaf is left in pageBuffer. + var leafPageId = FindLeafNodeWithPath(key, path, txnId, pageBuffer); var header = BTreeNodeHeader.ReadFrom(pageBuffer.AsSpan(32)); + if (_options.Unique) + { + var existing = BinarySearchLeaf(pageBuffer, header.EntryCount, key); + if (existing.Found) + { + if (!existing.Location.Equals(location)) + throw new InvalidOperationException("Duplicate key violation for unique index"); + return; + } + } + // Check if we need to split if (header.EntryCount >= MaxEntriesPerNode) { @@ -128,8 +129,7 @@ public void Insert(IndexKey key, DocumentLocation location, ulong? transactionId // Re-find leaf after split to ensure we have correct node path.Clear(); - leafPageId = FindLeafNodeWithPath(key, path, txnId); - ReadPage(leafPageId, txnId, pageBuffer); + leafPageId = FindLeafNodeWithPath(key, path, txnId, pageBuffer); } // Insert entry into leaf @@ -151,17 +151,24 @@ public bool TryFind(IndexKey key, out DocumentLocation location, ulong? transact var leafPageId = FindLeafNode(key, txnId); - Span pageBuffer = stackalloc byte[_storage.PageSize]; - ReadPage(leafPageId, txnId, pageBuffer); + var pageBuffer = System.Buffers.ArrayPool.Shared.Rent(_storage.PageSize); + try + { + ReadPage(leafPageId, txnId, pageBuffer); - var header = BTreeNodeHeader.ReadFrom(pageBuffer[32..]); - var result = BinarySearchLeaf(pageBuffer, header.EntryCount, key); - if (result.Found) + var header = BTreeNodeHeader.ReadFrom(pageBuffer.AsSpan(32)); + var result = BinarySearchLeaf(pageBuffer, header.EntryCount, key); + if (result.Found) + { + location = result.Location; + return true; + } + return false; + } + finally { - location = result.Location; - return true; + System.Buffers.ArrayPool.Shared.Return(pageBuffer); } - return false; } /// @@ -395,24 +402,10 @@ internal uint FindLeafNode(IndexKey key, ulong transactionId) private uint FindLeafNodeWithPath(IndexKey key, List path, ulong transactionId) { - var currentPageId = _rootPageId; var pageBuffer = System.Buffers.ArrayPool.Shared.Rent(_storage.PageSize); - try { - while (true) - { - ReadPage(currentPageId, transactionId, pageBuffer); - var header = BTreeNodeHeader.ReadFrom(pageBuffer.AsSpan(32)); - - if (header.IsLeaf) - { - return currentPageId; - } - - path.Add(currentPageId); - currentPageId = FindChildNode(pageBuffer, header, key); - } + return FindLeafNodeWithPath(key, path, transactionId, pageBuffer); } finally { @@ -420,6 +413,23 @@ private uint FindLeafNodeWithPath(IndexKey key, List path, ulong transacti } } + /// Like the overload above, but leaves the leaf page in . + private uint FindLeafNodeWithPath(IndexKey key, List path, ulong transactionId, byte[] pageBuffer) + { + var currentPageId = _rootPageId; + while (true) + { + ReadPage(currentPageId, transactionId, pageBuffer); + var header = BTreeNodeHeader.ReadFrom(pageBuffer.AsSpan(32)); + + if (header.IsLeaf) + return currentPageId; + + path.Add(currentPageId); + currentPageId = FindChildNode(pageBuffer, header, key); + } + } + private uint FindChildNode(Span nodeBuffer, BTreeNodeHeader header, IndexKey key) { // Internal Node Format: @@ -680,8 +690,7 @@ private void InsertIntoLeaf(uint leafPageId, IndexEntry entry, Span pageBu if (insertOffset < currentEnd) { int bytesToShift = currentEnd - insertOffset; - for (int i = bytesToShift - 1; i >= 0; i--) - pageBuffer[insertOffset + entrySize + i] = pageBuffer[insertOffset + i]; + pageBuffer.Slice(insertOffset, bytesToShift).CopyTo(pageBuffer.Slice(insertOffset + entrySize)); } // Write new entry at sorted insertion point diff --git a/src/BLite.Core/Query/BTreeExpressionVisitor.cs b/src/BLite.Core/Query/BTreeExpressionVisitor.cs index 3d556874..41f0eed4 100644 --- a/src/BLite.Core/Query/BTreeExpressionVisitor.cs +++ b/src/BLite.Core/Query/BTreeExpressionVisitor.cs @@ -44,6 +44,13 @@ protected override Expression VisitMethodCall(MethodCallExpression node) case "LongCount": VisitCount(node); return node; + case "First": + case "FirstOrDefault": + case "Single": + case "SingleOrDefault": + case "Any": + VisitElementTerminal(node); + return node; default: // GroupBy, Join, etc. — operators the direct pipeline cannot model. _model.HasComplexOperators = true; @@ -60,8 +67,11 @@ private void VisitWhere(MethodCallExpression node) Visit(node.Arguments[0]); var predicate = (UnaryExpression)node.Arguments[1]; - var lambda = (LambdaExpression)predicate.Operand; + AddPredicate((LambdaExpression)predicate.Operand); + } + private void AddPredicate(LambdaExpression lambda) + { if (_model.WhereClause == null) { _model.WhereClause = lambda; @@ -144,6 +154,33 @@ private void VisitAggregate(MethodCallExpression node) } } + /// + /// Handles First, FirstOrDefault, Single, SingleOrDefault and + /// Any, with or without a predicate. A predicate on the entities themselves is folded + /// into the WHERE clause (First(p) is Where(p).First()), so the index planner + /// sees it. Overloads taking a default value are left to the fallback pipeline. + /// + private void VisitElementTerminal(MethodCallExpression node) + { + Visit(node.Arguments[0]); // recurse into source + + if (node.Arguments.Count == 1) + { + _model.ElementTerminal = node.Method.Name; + } + else if (node.Arguments.Count == 2 + && node.Arguments[1] is UnaryExpression { NodeType: ExpressionType.Quote, Operand: LambdaExpression predicate } + && _model.SelectClause == null && !_model.Take.HasValue && !_model.Skip.HasValue) + { + AddPredicate(predicate); + _model.ElementTerminal = node.Method.Name; + } + else + { + _model.HasComplexOperators = true; + } + } + /// /// Handles Count() / LongCount() as a push-down scalar. /// When there is no predicate, the provider can use the primary BTree key count instead @@ -157,9 +194,15 @@ private void VisitCount(MethodCallExpression node) // Count() with no predicate — fast O(log n) + k BTree scan. _model.IsCountOnly = true; } + else if (_model.SelectClause == null && !_model.Take.HasValue && !_model.Skip.HasValue) + { + // Count(predicate) on the entities themselves is Where(predicate).Count(). + AddPredicate((LambdaExpression)((UnaryExpression)node.Arguments[1]).Operand); + _model.IsCountOnly = true; + } else { - // Count(predicate) — treat as a complex operator. + // Count(predicate) after a projection or Take/Skip — treat as a complex operator. _model.HasComplexOperators = true; } } diff --git a/src/BLite.Core/Query/BTreeQueryProvider.cs b/src/BLite.Core/Query/BTreeQueryProvider.cs index 32965ad9..dddaed2c 100644 --- a/src/BLite.Core/Query/BTreeQueryProvider.cs +++ b/src/BLite.Core/Query/BTreeQueryProvider.cs @@ -214,6 +214,33 @@ private async Task ExecuteAsyncCore(Expression expression, Can visitor.Visit(expression); var model = visitor.GetModel(); + // ── Element terminals: First / FirstOrDefault / Single / SingleOrDefault / Any ── + // With only a WHERE in front, at most one (Single*: two) matching document decides the + // result, so fetch just that many through the index-aware FetchAsync. Anything else in + // front (OrderBy, Select, Take, ...) goes through the general fallback as before. + if (model.ElementTerminal is { } terminal) + { + if (!model.HasComplexOperators && model.SelectClause == null && model.OrderByClause == null + && !model.Take.HasValue && !model.Skip.HasValue && model.AggregateOp == null && !model.IsCountOnly) + { + bool single = terminal is "Single" or "SingleOrDefault"; + var found = new List(single ? 2 : 1); + await foreach (var item in _collection.FetchAsync(model.WhereClause, single ? 2 : 1, cancellationToken)) + found.Add(item); + + return terminal switch + { + "First" => (TResult)(object)found.First()!, + "FirstOrDefault" => (TResult)(object?)found.FirstOrDefault()!, + "Single" => (TResult)(object)found.Single()!, + "SingleOrDefault" => (TResult)(object?)found.SingleOrDefault()!, + _ => (TResult)(object)(found.Count > 0), + }; + } + + model.HasComplexOperators = true; + } + // ── Fast path: Count() / LongCount() terminal with no WHERE ───────────── // Uses the primary BTree key scan (O(n) key reads, zero document reads). // Disabled when Take/Skip are present — those shaping operators reduce the diff --git a/src/BLite.Core/Query/QueryModel.cs b/src/BLite.Core/Query/QueryModel.cs index e998fadd..7bb97fb7 100644 --- a/src/BLite.Core/Query/QueryModel.cs +++ b/src/BLite.Core/Query/QueryModel.cs @@ -34,4 +34,11 @@ internal class QueryModel /// Enables an O(n) primary-BTree key scan instead of full document deserialisation. /// public bool IsCountOnly { get; set; } + + /// + /// Element terminal (First, FirstOrDefault, Single, + /// SingleOrDefault or Any) ending the query, or null. A predicate passed to it + /// has already been folded into . + /// + public string? ElementTerminal { get; set; } } diff --git a/src/BLite.Core/Transactions/Transaction.cs b/src/BLite.Core/Transactions/Transaction.cs index 4dfd990b..519e09f0 100644 --- a/src/BLite.Core/Transactions/Transaction.cs +++ b/src/BLite.Core/Transactions/Transaction.cs @@ -58,10 +58,9 @@ public void AddWrite(WriteOperation operation) if (_state != TransactionState.Active) throw new InvalidOperationException($"Cannot add writes to transaction in state {_state}"); - // Defensive copy: necessary to prevent use-after-return if caller uses pooled buffers - byte[] ownedCopy = operation.NewValue.ToArray(); - // StorageEngine gestisce tutte le scritture transazionali - _storage.WritePage(operation.PageId, _transactionId, ownedCopy); + // StorageEngine.WritePage copies the page into the transaction's own buffer, so a caller + // may return a pooled buffer as soon as this call returns. + _storage.WritePage(operation.PageId, _transactionId, operation.NewValue.Span); } /// diff --git a/tests/BLite.Tests/QueryAndWritePathTests.cs b/tests/BLite.Tests/QueryAndWritePathTests.cs new file mode 100644 index 00000000..116497e6 --- /dev/null +++ b/tests/BLite.Tests/QueryAndWritePathTests.cs @@ -0,0 +1,167 @@ +using System.Linq.Expressions; +using BLite.Bson; +using BLite.Core; +using BLite.Core.Indexing; +using BLite.Core.Query; +using BLite.Core.Storage; +using BLite.Core.Transactions; +using BLite.Shared; + +namespace BLite.Tests; + +/// +/// Predicate overloads of LINQ terminals use the index planner, and the write path changes +/// (single-pass B-Tree insert, no defensive page copy) keep their semantics. +/// +public class QueryAndWritePathTests : IDisposable +{ + private readonly string _dbPath; + private readonly TestDbContext _db; + private readonly List _people; + + public QueryAndWritePathTests() + { + _dbPath = Path.Combine(Path.GetTempPath(), $"query_write_{Guid.NewGuid():N}.db"); + _db = new TestDbContext(_dbPath); + _people = Enumerable.Range(0, 300).Select(i => new PersonV2 { Name = $"p{i}", Age = i % 50 }).ToList(); + _db.PeopleV2.InsertBulkAsync(_people).GetAwaiter().GetResult(); + } + + public void Dispose() + { + _db.Dispose(); + foreach (var f in new[] { _dbPath, Path.ChangeExtension(_dbPath, ".wal") }) + try { if (File.Exists(f)) File.Delete(f); } catch { /* best-effort */ } + } + + private static QueryModel Model(Expression expression) + { + var visitor = new BTreeExpressionVisitor(); + visitor.Visit(expression); + return visitor.GetModel(); + } + + private IQueryable Q => _db.PeopleV2.AsQueryable(); + + // ── Query model ───────────────────────────────────────────────────────── + + [Fact] + public void Model_PredicateOfCountAndElementTerminals_FoldsIntoWhere() + { + var count = Model(Expression.Call(typeof(Queryable), "Count", [typeof(PersonV2)], Q.Expression, Expression.Quote((Expression>)(p => p.Age == 3)))); + Assert.True(count.IsCountOnly); + Assert.False(count.HasComplexOperators); + Assert.NotNull(count.WhereClause); + + foreach (var name in new[] { "First", "FirstOrDefault", "Single", "SingleOrDefault", "Any" }) + { + var model = Model(Expression.Call(typeof(Queryable), name, [typeof(PersonV2)], Q.Expression, Expression.Quote((Expression>)(p => p.Age == 3)))); + Assert.Equal(name, model.ElementTerminal); + Assert.False(model.HasComplexOperators); + Assert.NotNull(model.WhereClause); + } + } + + [Fact] + public void Model_PredicateAfterTakeSkipOrSelect_IsNotFolded() + { + var afterTake = Model(Expression.Call(typeof(Queryable), "Count", [typeof(PersonV2)], Q.Take(5).Expression, Expression.Quote((Expression>)(p => p.Age == 3)))); + var afterSkip = Model(Expression.Call(typeof(Queryable), "First", [typeof(PersonV2)], Q.Skip(5).Expression, Expression.Quote((Expression>)(p => p.Age == 3)))); + var afterSelect = Model(Expression.Call(typeof(Queryable), "Any", [typeof(int)], Q.Select(p => p.Age).Expression, Expression.Quote((Expression>)(a => a == 3)))); + + Assert.True(afterTake.HasComplexOperators); + Assert.Null(afterTake.WhereClause); + Assert.True(afterSkip.HasComplexOperators); + Assert.Null(afterSkip.WhereClause); + Assert.True(afterSelect.HasComplexOperators); + Assert.Null(afterSelect.WhereClause); + } + + // ── Results ───────────────────────────────────────────────────────────── + + [Fact] + public void Count_WithPredicate_MatchesLinqToObjects() + { + Assert.Equal(_people.Count(p => p.Age == 7), Q.Count(p => p.Age == 7)); + Assert.Equal(_people.Count(p => p.Age >= 10 && p.Age < 20), Q.Count(p => p.Age >= 10 && p.Age < 20)); + Assert.Equal(_people.Count(p => p.Name.EndsWith("9")), Q.Count(p => p.Name.EndsWith("9"))); + Assert.Equal(_people.LongCount(p => p.Age == 7), Q.LongCount(p => p.Age == 7)); + Assert.Equal(_people.OrderBy(p => p.Name).Take(40).Count(p => p.Age == 7), Q.OrderBy(p => p.Name).Take(40).Count(p => p.Age == 7)); + } + + [Fact] + public void ElementTerminals_WithPredicate_MatchLinqToObjects() + { + Assert.Equal(7, Q.First(p => p.Age == 7).Age); + Assert.Equal("p42", Q.FirstOrDefault(p => p.Name == "p42")!.Name); + Assert.Null(Q.FirstOrDefault(p => p.Age == 999)); + Assert.Equal("p42", Q.Single(p => p.Name == "p42").Name); + Assert.Null(Q.SingleOrDefault(p => p.Name == "nobody")); + Assert.True(Q.Any(p => p.Age == 49)); + Assert.False(Q.Any(p => p.Age == 50)); + Assert.True(Q.Any()); + + Assert.Throws(() => Q.First(p => p.Age == 999)); + Assert.Throws(() => Q.Single(p => p.Age == 7)); + Assert.Throws(() => Q.SingleOrDefault(p => p.Age == 7)); + Assert.Throws(() => Q.Single(p => p.Age == 999)); + } + + [Fact] + public void ElementTerminals_AfterOtherOperators_KeepTheirSemantics() + { + Assert.Equal(_people.OrderByDescending(p => p.Name).First(p => p.Age == 7).Name, + Q.OrderByDescending(p => p.Name).First(p => p.Age == 7).Name); + Assert.Equal(_people.Where(p => p.Age > 5).FirstOrDefault(p => p.Age < 8)!.Age, + Q.Where(p => p.Age > 5).FirstOrDefault(p => p.Age < 8)!.Age); + Assert.Equal(_people.OrderBy(p => p.Name).Skip(290).Any(p => p.Age == 7), Q.OrderBy(p => p.Name).Skip(290).Any(p => p.Age == 7)); + Assert.True(Q.All(p => p.Age < 50)); + + var fallback = new PersonV2 { Name = "fallback" }; + Assert.Same(fallback, Q.FirstOrDefault(p => p.Age == 999, fallback)); + } + + // ── Write path ────────────────────────────────────────────────────────── + + [Fact] + public async Task AddWrite_CallerMayReuseItsBufferRightAway() + { + var page = _db.Storage.AllocatePage(); + var buffer = new byte[_db.Storage.PageSize]; + buffer.AsSpan(100, 4).Fill(0xAB); + + using (var txn = _db.Storage.BeginTransaction()) + { + txn.AddWrite(new WriteOperation(ObjectId.Empty, buffer.AsMemory(), page, OperationType.Insert)); + buffer.AsSpan().Clear(); // e.g. returned to a pool and reused + + var read = new byte[_db.Storage.PageSize]; + _db.Storage.ReadPage(page, txn.TransactionId, read); + Assert.All(read.AsSpan(100, 4).ToArray(), b => Assert.Equal(0xAB, b)); + await txn.CommitAsync(); + } + + var committed = new byte[_db.Storage.PageSize]; + _db.Storage.ReadPage(page, null, committed); + Assert.All(committed.AsSpan(100, 4).ToArray(), b => Assert.Equal(0xAB, b)); + } + + [Fact] + public async Task UniqueBTree_SinglePassInsert_RejectsDuplicatesAndIgnoresExactReinserts() + { + var tree = new BTreeIndex(_db.Storage, IndexOptions.CreateUnique("_id")); + var keys = Enumerable.Range(0, 500).Select(_ => new IndexKey(ObjectId.NewObjectId())).ToArray(); + + using var txn = _db.Storage.BeginTransaction(); + for (int i = 0; i < keys.Length; i++) + tree.Insert(keys[i], new DocumentLocation((uint)(i + 1), 0), txn.TransactionId); + + tree.Insert(keys[123], new DocumentLocation(124, 0), txn.TransactionId); // same key and location: no-op + Assert.Throws(() => tree.Insert(keys[123], new DocumentLocation(999, 0), txn.TransactionId)); + + Assert.Equal(500, tree.Range(IndexKey.MinKey, IndexKey.MaxKey, IndexDirection.Forward, txn.TransactionId).Count()); + Assert.True(tree.TryFind(keys[123], out var location, txn.TransactionId)); + Assert.Equal(124u, location.PageId); + await txn.CommitAsync(); + } +} From ef05b293ffa47623fd6c5179d243fa4ba8677984 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ricardo=20Gro=C3=9F?= Date: Mon, 12 Oct 2026 17:07:20 +0200 Subject: [PATCH 2/4] fix(query): fold terminal predicates only after Where/OrderBy/ThenBy (#180) Address review of #180: - Count/LongCount/First*/Single*/Any fold their predicate into the WHERE clause only when every operator between the collection root and the terminal is a single-parameter Where or OrderBy*/ThenBy*, and the predicate takes the collection's entity type. Fixes the InvalidCastException after GroupBy/OfType/Cast/SelectMany/Join/Zip and the wrong results after SkipWhile/TakeWhile (and the indexed Where overload). - Regression tests for GroupBy (Count/Any/First), SkipWhile, TakeWhile, OfType/Cast, the fold allowlist, and Count(p => p.V.VectorSearch(q, k)). - Transaction.AddWrite XML doc no longer claims a defensive copy. - AllocateOverflowPage passes the pooled buffer to AddWrite without a redundant ToArray(). Co-Authored-By: Claude Opus 5.5 --- .../Collections/DocumentCollection.cs | 2 +- .../Query/BTreeExpressionVisitor.cs | 51 ++++++++++++-- src/BLite.Core/Transactions/Transaction.cs | 5 +- tests/BLite.Tests/QueryAndWritePathTests.cs | 68 +++++++++++++++++++ 4 files changed, 118 insertions(+), 8 deletions(-) diff --git a/src/BLite.Core/Collections/DocumentCollection.cs b/src/BLite.Core/Collections/DocumentCollection.cs index 58934190..605ed62d 100644 --- a/src/BLite.Core/Collections/DocumentCollection.cs +++ b/src/BLite.Core/Collections/DocumentCollection.cs @@ -1920,7 +1920,7 @@ private uint AllocateOverflowPage(ReadOnlySpan data, uint nextOverflowPage // NEW: Buffer write in transaction or write immediately var writeOp = new WriteOperation( documentId: ObjectId.Empty, - newValue: buffer.AsSpan(0, _storage.PageSize).ToArray(), + newValue: buffer.AsMemory(0, _storage.PageSize), pageId: pageId, type: OperationType.Insert ); diff --git a/src/BLite.Core/Query/BTreeExpressionVisitor.cs b/src/BLite.Core/Query/BTreeExpressionVisitor.cs index 41f0eed4..59e8f64f 100644 --- a/src/BLite.Core/Query/BTreeExpressionVisitor.cs +++ b/src/BLite.Core/Query/BTreeExpressionVisitor.cs @@ -158,7 +158,8 @@ private void VisitAggregate(MethodCallExpression node) /// Handles First, FirstOrDefault, Single, SingleOrDefault and /// Any, with or without a predicate. A predicate on the entities themselves is folded /// into the WHERE clause (First(p) is Where(p).First()), so the index planner - /// sees it. Overloads taking a default value are left to the fallback pipeline. + /// sees it, but only where allows it. Overloads taking a + /// default value are left to the fallback pipeline. /// private void VisitElementTerminal(MethodCallExpression node) { @@ -170,7 +171,7 @@ private void VisitElementTerminal(MethodCallExpression node) } else if (node.Arguments.Count == 2 && node.Arguments[1] is UnaryExpression { NodeType: ExpressionType.Quote, Operand: LambdaExpression predicate } - && _model.SelectClause == null && !_model.Take.HasValue && !_model.Skip.HasValue) + && CanFoldPredicate(node.Arguments[0], predicate)) { AddPredicate(predicate); _model.ElementTerminal = node.Method.Name; @@ -194,16 +195,56 @@ private void VisitCount(MethodCallExpression node) // Count() with no predicate — fast O(log n) + k BTree scan. _model.IsCountOnly = true; } - else if (_model.SelectClause == null && !_model.Take.HasValue && !_model.Skip.HasValue) + else if (node.Arguments[1] is UnaryExpression { NodeType: ExpressionType.Quote, Operand: LambdaExpression predicate } + && CanFoldPredicate(node.Arguments[0], predicate)) { // Count(predicate) on the entities themselves is Where(predicate).Count(). - AddPredicate((LambdaExpression)((UnaryExpression)node.Arguments[1]).Operand); + AddPredicate(predicate); _model.IsCountOnly = true; } else { - // Count(predicate) after a projection or Take/Skip — treat as a complex operator. + // Count(predicate) after any other operator — treat as a complex operator. _model.HasComplexOperators = true; } } + + /// + /// True when a terminal's may be folded into the WHERE clause, + /// i.e. when source.Terminal(p) equals source.Where(p).Terminal() evaluated with + /// the filter applied first. That holds only when every operator between the root and the + /// terminal is a plain Where or an ordering (OrderBy* / ThenBy*): those keep + /// the element type and commute with a filter. Anything else, such as GroupBy, + /// OfType, Select or the order-dependent Take, Skip, + /// SkipWhile and TakeWhile, either changes what the predicate is applied to or + /// gives a different result when the predicate runs first. + /// + private static bool CanFoldPredicate(Expression source, LambdaExpression predicate) + { + var current = source; + while (current is MethodCallExpression call) + { + if (call.Method.DeclaringType != typeof(Queryable)) + return false; + + switch (call.Method.Name) + { + case "Where" when call.Arguments[1] is UnaryExpression { Operand: LambdaExpression { Parameters.Count: 1 } }: + case "OrderBy": + case "OrderByDescending": + case "ThenBy": + case "ThenByDescending": + current = call.Arguments[0]; + break; + default: + return false; + } + } + + // Only element-type-preserving operators were skipped, so the predicate must take the + // root collection's own element type (the entity type). + return current is ConstantExpression { Value: IQueryable root } + && predicate.Parameters.Count == 1 + && predicate.Parameters[0].Type == root.ElementType; + } } diff --git a/src/BLite.Core/Transactions/Transaction.cs b/src/BLite.Core/Transactions/Transaction.cs index 519e09f0..41d6b722 100644 --- a/src/BLite.Core/Transactions/Transaction.cs +++ b/src/BLite.Core/Transactions/Transaction.cs @@ -50,8 +50,9 @@ internal void AddChange(CDC.InternalChangeEvent change) /// /// Adds a write operation to the transaction's write set. - /// NOTE: Makes a defensive copy of the data to ensure memory safety. - /// This allocation is necessary because the caller may return the buffer to a pool. + /// The page is copied into the transaction's own buffer by , + /// so only has to stay valid for the duration of this call + /// and the caller may return a pooled buffer right afterwards. /// public void AddWrite(WriteOperation operation) { diff --git a/tests/BLite.Tests/QueryAndWritePathTests.cs b/tests/BLite.Tests/QueryAndWritePathTests.cs index 116497e6..d0034d2e 100644 --- a/tests/BLite.Tests/QueryAndWritePathTests.cs +++ b/tests/BLite.Tests/QueryAndWritePathTests.cs @@ -121,6 +121,74 @@ public void ElementTerminals_AfterOtherOperators_KeepTheirSemantics() Assert.Same(fallback, Q.FirstOrDefault(p => p.Age == 999, fallback)); } + [Fact] + public void Terminals_AfterGroupBy_PredicateIsOverGroups() + { + // The predicate takes an IGrouping, not a PersonV2, so it must not reach FetchAsync. + Assert.Equal(_people.GroupBy(p => p.Age).Count(g => g.Count() > 5), Q.GroupBy(p => p.Age).Count(g => g.Count() > 5)); + Assert.True(Q.GroupBy(p => p.Age).Any(g => g.Key == 3)); + Assert.Equal(3, Q.GroupBy(p => p.Age).First(g => g.Key == 3).Key); + } + + [Fact] + public void Terminals_AfterSkipWhile_PredicateIsNotPushedBelowIt() + { + // SkipWhile/TakeWhile depend on the order the collection yields its documents in. + var scan = Q.ToList(); + Assert.Equal(scan.SkipWhile(p => p.Age < 10).FirstOrDefault(p => p.Age == 5)?.Name, + Q.SkipWhile(p => p.Age < 10).FirstOrDefault(p => p.Age == 5)?.Name); + Assert.Equal(scan.SkipWhile(p => p.Age < 10).Count(p => p.Age == 5), + Q.SkipWhile(p => p.Age < 10).Count(p => p.Age == 5)); + } + + [Fact] + public void Terminals_AfterTakeWhile_PredicateIsNotPushedBelowIt() + { + var scan = Q.ToList(); + Assert.Equal(scan.TakeWhile(p => p.Age < 10).Count(p => p.Age == 5), + Q.TakeWhile(p => p.Age < 10).Count(p => p.Age == 5)); + Assert.Equal(scan.TakeWhile(p => p.Age < 10).Any(p => p.Age == 5), + Q.TakeWhile(p => p.Age < 10).Any(p => p.Age == 5)); + } + + [Fact] + public void Terminals_AfterOfTypeOrCast_PredicateIsOverTheNewElementType() + { + Assert.Equal(_people.Count, Q.OfType().Count(o => o != null)); + Assert.True(Q.Cast().Any(o => o is PersonV2)); + Assert.Equal(_people.Count(p => p.Age == 3), Q.OfType().Count(p => p.Age == 3)); + } + + [Fact] + public void Model_PredicateIsFoldedOnlyAfterWhereAndOrdering() + { + Expression> pred = p => p.Age == 3; + + var afterOrdering = Model(Expression.Call(typeof(Queryable), "First", [typeof(PersonV2)], + Q.Where(p => p.Age > 1).OrderBy(p => p.Name).ThenByDescending(p => p.Age).Expression, Expression.Quote(pred))); + Assert.NotNull(afterOrdering.WhereClause); + Assert.Equal("First", afterOrdering.ElementTerminal); + + var afterSkipWhile = Model(Expression.Call(typeof(Queryable), "Count", [typeof(PersonV2)], + Q.SkipWhile(p => p.Age < 10).Expression, Expression.Quote(pred))); + Assert.True(afterSkipWhile.HasComplexOperators); + Assert.Null(afterSkipWhile.WhereClause); + + var afterIndexedWhere = Model(Expression.Call(typeof(Queryable), "Any", [typeof(PersonV2)], + Q.Where((p, i) => i > 10).Expression, Expression.Quote(pred))); + Assert.True(afterIndexedWhere.HasComplexOperators); + } + + [Fact] + public async Task Count_WithVectorSearchPredicate_CountsTheNearestK() + { + for (int i = 0; i < 20; i++) + await _db.VectorItems.InsertAsync(new VectorEntity { Title = $"v{i}", Embedding = [i, 0f, 0f] }); + + float[] query = [0f, 0f, 0f]; + Assert.Equal(3, _db.VectorItems.AsQueryable().Count(x => x.Embedding.VectorSearch(query, 3))); + } + // ── Write path ────────────────────────────────────────────────────────── [Fact] From a167c72f90b71037a4efafd8baf2b79ce9fa6236 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ricardo=20Gro=C3=9F?= Date: Tue, 13 Oct 2026 15:13:26 +0200 Subject: [PATCH 3/4] test(query): expect Count(VectorSearch) after Where to be folded (#180) With #173 on main, Where(x => x.Title != null).Count(x => x.Embedding.VectorSearch(Q, 1)) has its predicate folded into the WHERE, so the index serves the VectorSearch instead of the replay throwing. - Rename VectorSearch_InCountPredicate_AfterOtherWhere_Throws to VectorSearch_InCountPredicate_AfterOtherWhere_IsFoldedIntoTheWhere and assert the same count as Where(x => x.Title != null && x.Embedding.VectorSearch(Q, 1)).Count(). - VectorSearch_InSumSelector_AfterOtherWhere_Throws (not folded) is unchanged. Co-Authored-By: Claude Opus 5.5 --- .../BLite.Tests/NearWithinCorrectnessTests.cs | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tests/BLite.Tests/NearWithinCorrectnessTests.cs b/tests/BLite.Tests/NearWithinCorrectnessTests.cs index 9fc5ae69..6721a4ea 100644 --- a/tests/BLite.Tests/NearWithinCorrectnessTests.cs +++ b/tests/BLite.Tests/NearWithinCorrectnessTests.cs @@ -199,8 +199,9 @@ public async Task TwoVectorSearches_ExplainThatOnlyOneIsSupported() Assert.Contains("only one VectorSearch", ex.Message); } - // Only the VectorSearch in a Where that FetchAsync applied is served by the index. One in a - // projection, a group or an aggregate after the fetch must still throw, not match everything. + // Only the VectorSearch in a Where that FetchAsync applied (or in a terminal predicate folded + // into it) is served by the index. One in a projection, a group or an aggregate selector after + // the fetch must still throw, not match everything. private static readonly float[] Q = [0f, 0f, 0f]; @@ -272,14 +273,22 @@ public async Task VectorSearch_InSumSelector_AfterOtherWhere_Throws() .Sum(x => x.Embedding.VectorSearch(Q, 1) ? 1 : 0)); } + // Count(p) after a plain Where is folded into the WHERE (#180), so its VectorSearch is served + // by the index like any other: the vector term drives and Title != null is the post-filter. [Fact] - public async Task VectorSearch_InCountPredicate_AfterOtherWhere_Throws() + public async Task VectorSearch_InCountPredicate_AfterOtherWhere_IsFoldedIntoTheWhere() { await InsertVectorProbeDataAsync(); - Assert.Throws(() => _db.VectorItems.AsQueryable() + var folded = _db.VectorItems.AsQueryable() .Where(x => x.Title != null) - .Count(x => x.Embedding.VectorSearch(Q, 1))); + .Count(x => x.Embedding.VectorSearch(Q, 1)); + var combined = _db.VectorItems.AsQueryable() + .Where(x => x.Title != null && x.Embedding.VectorSearch(Q, 1)) + .Count(); + + Assert.Equal(combined, folded); + Assert.Equal(1, folded); } // ── DynamicCollection ─────────────────────────────────────────────────── From 3cba537ff86be0a45f15daec126b6d6353c0a109 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ricardo=20Gro=C3=9F?= Date: Tue, 13 Oct 2026 15:20:40 +0200 Subject: [PATCH 4/4] fix(query): don't replay a terminal predicate folded into the WHERE (#180) OrderBy(x => x.Title).First(x => x.Embedding.VectorSearch(Q, 1)) folded the predicate into the WHERE, but the OrderBy sent the element terminal through the fallback, whose replay evaluated First(p) again and threw for VectorSearch. The same happened for Any, and for Count/LongCount after ThenBy (which the direct pipeline does not model). - QueryModel.FoldedTerminal records the terminal whose predicate was folded; the replay swaps it for its parameterless overload (First(p) -> First()) next to DropAppliedWheres. FetchAsync already applied p, and folding requires only Where/OrderBy/ThenBy in front. - Last/LastOrDefault fold their predicate too (Last(p) is Where(p).Last()); they always use the fallback since they need every match. - All(p) is still never folded and still evaluated in the replay, as are predicates after operators that block folding (SkipWhile, Take, GroupBy, ...). - Tests: OrderBy/ThenBy + First/FirstOrDefault/Last/LastOrDefault/Single/SingleOrDefault/ Any/Count/LongCount with VectorSearch against the equivalent Where(...).OrderBy(...) query, the same terminals with a plain predicate against LINQ-to-Objects, and unfolded shapes (All, SkipWhile, Take, GroupBy) that still throw. Co-Authored-By: Claude Opus 5.5 --- .../Query/BTreeExpressionVisitor.cs | 15 ++- src/BLite.Core/Query/BTreeQueryProvider.cs | 37 +++++- src/BLite.Core/Query/QueryModel.cs | 10 +- tests/BLite.Tests/QueryAndWritePathTests.cs | 106 +++++++++++++++++- 4 files changed, 157 insertions(+), 11 deletions(-) diff --git a/src/BLite.Core/Query/BTreeExpressionVisitor.cs b/src/BLite.Core/Query/BTreeExpressionVisitor.cs index 59e8f64f..43866a2c 100644 --- a/src/BLite.Core/Query/BTreeExpressionVisitor.cs +++ b/src/BLite.Core/Query/BTreeExpressionVisitor.cs @@ -48,6 +48,8 @@ protected override Expression VisitMethodCall(MethodCallExpression node) case "FirstOrDefault": case "Single": case "SingleOrDefault": + case "Last": + case "LastOrDefault": case "Any": VisitElementTerminal(node); return node; @@ -155,11 +157,12 @@ private void VisitAggregate(MethodCallExpression node) } /// - /// Handles First, FirstOrDefault, Single, SingleOrDefault and - /// Any, with or without a predicate. A predicate on the entities themselves is folded - /// into the WHERE clause (First(p) is Where(p).First()), so the index planner - /// sees it, but only where allows it. Overloads taking a - /// default value are left to the fallback pipeline. + /// Handles First, FirstOrDefault, Single, SingleOrDefault, + /// Last, LastOrDefault and Any, with or without a predicate. A predicate + /// on the entities themselves is folded into the WHERE clause (First(p) is + /// Where(p).First()), so the index planner sees it, but only where + /// allows it. Overloads taking a default value are left to the + /// fallback pipeline. All(p) is not handled here: it is not Where(p).All(). /// private void VisitElementTerminal(MethodCallExpression node) { @@ -175,6 +178,7 @@ private void VisitElementTerminal(MethodCallExpression node) { AddPredicate(predicate); _model.ElementTerminal = node.Method.Name; + _model.FoldedTerminal = node; } else { @@ -201,6 +205,7 @@ private void VisitCount(MethodCallExpression node) // Count(predicate) on the entities themselves is Where(predicate).Count(). AddPredicate(predicate); _model.IsCountOnly = true; + _model.FoldedTerminal = node; } else { diff --git a/src/BLite.Core/Query/BTreeQueryProvider.cs b/src/BLite.Core/Query/BTreeQueryProvider.cs index dddaed2c..518e75c4 100644 --- a/src/BLite.Core/Query/BTreeQueryProvider.cs +++ b/src/BLite.Core/Query/BTreeQueryProvider.cs @@ -217,10 +217,12 @@ private async Task ExecuteAsyncCore(Expression expression, Can // ── Element terminals: First / FirstOrDefault / Single / SingleOrDefault / Any ── // With only a WHERE in front, at most one (Single*: two) matching document decides the // result, so fetch just that many through the index-aware FetchAsync. Anything else in - // front (OrderBy, Select, Take, ...) goes through the general fallback as before. + // front (OrderBy, Select, Take, ...), and Last / LastOrDefault, which need every match, + // go through the general fallback. if (model.ElementTerminal is { } terminal) { - if (!model.HasComplexOperators && model.SelectClause == null && model.OrderByClause == null + if (terminal is not ("Last" or "LastOrDefault") + && !model.HasComplexOperators && model.SelectClause == null && model.OrderByClause == null && !model.Take.HasValue && !model.Skip.HasValue && model.AggregateOp == null && !model.IsCountOnly) { bool single = terminal is "Single" or "SingleOrDefault"; @@ -341,8 +343,9 @@ private async Task ExecuteAsyncCore(Expression expression, Can // FetchAsync applied the WHERE, so the replay drops the root-chain Where calls folded into // it. Re-evaluating them would change nothing except for VectorSearch, which only the index // can evaluate. Every other VectorSearch in the expression is still evaluated and throws. + // The same holds for a terminal predicate folded into the WHERE (First(p) after OrderBy). if (whereAlreadyApplied) - expression = DropAppliedWheres(expression).Expression; + expression = DropAppliedWheres(DropFoldedTerminalPredicate(expression, model.FoldedTerminal)).Expression; if (model.HasComplexOperators) return ExecuteViaEnumerableRewriter(expression, sourceData); @@ -702,6 +705,34 @@ private static (Expression Expression, bool IsFetchedSubset) DropAppliedWheres(E return (rewritten, isFetchedSubset && s_subsetPreservingOperators.Contains(call.Method.Name)); } + /// + /// Replaces the terminal whose predicate folded into + /// (Count(p), First(p), ...) with its + /// parameterless overload. The visitor folds a predicate only when the terminal's source is the + /// root followed by plain Where and ordering calls, so the elements reaching the terminal are + /// fetched documents, which all satisfy the predicate already. + /// + [RequiresDynamicCode("BLite LINQ queries use Expression.Compile() and MakeGenericMethod which require dynamic code generation.")] + [RequiresUnreferencedCode("BLite LINQ queries use reflection to resolve methods and types at runtime. Ensure all entity types and their members are preserved.")] + private static Expression DropFoldedTerminalPredicate(Expression expression, MethodCallExpression? foldedTerminal) + { + if (foldedTerminal is null || !ReferenceEquals(expression, foldedTerminal)) + return expression; + + var name = foldedTerminal.Method.Name; + if (!s_parameterlessTerminals.TryGetValue(name, out var definition)) + { + definition = typeof(Queryable).GetMethods().Single(m => m.Name == name && m.GetParameters().Length == 1); + s_parameterlessTerminals[name] = definition; + } + + var parameterless = definition.MakeGenericMethod(foldedTerminal.Method.GetGenericArguments()); + return Expression.Call(parameterless, foldedTerminal.Arguments[0]); + } + + /// Generic definitions of the parameterless Queryable terminals, by name. + private static readonly ConcurrentDictionary s_parameterlessTerminals = new(); + /// Queryable operators whose output elements are a subset of their source's. private static readonly HashSet s_subsetPreservingOperators = [ diff --git a/src/BLite.Core/Query/QueryModel.cs b/src/BLite.Core/Query/QueryModel.cs index 7bb97fb7..ed06955c 100644 --- a/src/BLite.Core/Query/QueryModel.cs +++ b/src/BLite.Core/Query/QueryModel.cs @@ -37,8 +37,14 @@ internal class QueryModel /// /// Element terminal (First, FirstOrDefault, Single, - /// SingleOrDefault or Any) ending the query, or null. A predicate passed to it - /// has already been folded into . + /// SingleOrDefault, Last, LastOrDefault or Any) ending the query, + /// or null. A predicate passed to it has already been folded into . /// public string? ElementTerminal { get; set; } + + /// + /// The terminal call (Count(p), First(p), ...) whose predicate was folded into + /// , or null. A fallback replay must not evaluate that predicate again. + /// + public MethodCallExpression? FoldedTerminal { get; set; } } diff --git a/tests/BLite.Tests/QueryAndWritePathTests.cs b/tests/BLite.Tests/QueryAndWritePathTests.cs index d0034d2e..8be51ad3 100644 --- a/tests/BLite.Tests/QueryAndWritePathTests.cs +++ b/tests/BLite.Tests/QueryAndWritePathTests.cs @@ -53,15 +53,29 @@ public void Model_PredicateOfCountAndElementTerminals_FoldsIntoWhere() Assert.False(count.HasComplexOperators); Assert.NotNull(count.WhereClause); - foreach (var name in new[] { "First", "FirstOrDefault", "Single", "SingleOrDefault", "Any" }) + Assert.NotNull(count.FoldedTerminal); + + foreach (var name in new[] { "First", "FirstOrDefault", "Single", "SingleOrDefault", "Last", "LastOrDefault", "Any" }) { var model = Model(Expression.Call(typeof(Queryable), name, [typeof(PersonV2)], Q.Expression, Expression.Quote((Expression>)(p => p.Age == 3)))); Assert.Equal(name, model.ElementTerminal); Assert.False(model.HasComplexOperators); Assert.NotNull(model.WhereClause); + Assert.NotNull(model.FoldedTerminal); } } + [Fact] + public void Model_AllPredicate_IsNeverFolded() + { + // All(p) is not Where(p).All(): it must see the documents that fail p. + var all = Model(Expression.Call(typeof(Queryable), "All", [typeof(PersonV2)], + Q.OrderBy(p => p.Name).Expression, Expression.Quote((Expression>)(p => p.Age == 3)))); + Assert.True(all.HasComplexOperators); + Assert.Null(all.WhereClause); + Assert.Null(all.FoldedTerminal); + } + [Fact] public void Model_PredicateAfterTakeSkipOrSelect_IsNotFolded() { @@ -173,6 +187,7 @@ public void Model_PredicateIsFoldedOnlyAfterWhereAndOrdering() Q.SkipWhile(p => p.Age < 10).Expression, Expression.Quote(pred))); Assert.True(afterSkipWhile.HasComplexOperators); Assert.Null(afterSkipWhile.WhereClause); + Assert.Null(afterSkipWhile.FoldedTerminal); var afterIndexedWhere = Model(Expression.Call(typeof(Queryable), "Any", [typeof(PersonV2)], Q.Where((p, i) => i > 10).Expression, Expression.Quote(pred))); @@ -189,6 +204,95 @@ public async Task Count_WithVectorSearchPredicate_CountsTheNearestK() Assert.Equal(3, _db.VectorItems.AsQueryable().Count(x => x.Embedding.VectorSearch(query, 3))); } + [Fact] + public void Terminals_AfterOrderingWithPredicate_MatchLinqToObjects() + { + // The predicate is folded into the WHERE and dropped from the replay, but the ordering + // must still decide which matching document the terminal returns. + Assert.Equal(_people.OrderBy(p => p.Name).First(p => p.Age == 7).Name, + Q.OrderBy(p => p.Name).First(p => p.Age == 7).Name); + Assert.Equal(_people.OrderBy(p => p.Name).FirstOrDefault(p => p.Age == 7)!.Name, + Q.OrderBy(p => p.Name).FirstOrDefault(p => p.Age == 7)!.Name); + Assert.Null(Q.OrderBy(p => p.Name).FirstOrDefault(p => p.Age == 999)); + Assert.Equal(_people.OrderBy(p => p.Name).Last(p => p.Age == 7).Name, + Q.OrderBy(p => p.Name).Last(p => p.Age == 7).Name); + Assert.Equal(_people.OrderByDescending(p => p.Name).LastOrDefault(p => p.Age == 7)!.Name, + Q.OrderByDescending(p => p.Name).LastOrDefault(p => p.Age == 7)!.Name); + Assert.Null(Q.OrderBy(p => p.Name).LastOrDefault(p => p.Age == 999)); + Assert.Equal("p42", Q.OrderBy(p => p.Name).Single(p => p.Name == "p42").Name); + Assert.Null(Q.OrderBy(p => p.Name).SingleOrDefault(p => p.Name == "nobody")); + Assert.Throws(() => Q.OrderBy(p => p.Name).Single(p => p.Age == 7)); + Assert.True(Q.OrderBy(p => p.Name).Any(p => p.Age == 49)); + Assert.False(Q.OrderBy(p => p.Name).Any(p => p.Age == 50)); + + Assert.Equal(_people.OrderBy(p => p.Age).ThenByDescending(p => p.Name).First(p => p.Name.EndsWith("3")).Name, + Q.OrderBy(p => p.Age).ThenByDescending(p => p.Name).First(p => p.Name.EndsWith("3")).Name); + Assert.Equal(_people.Where(p => p.Age > 10).OrderBy(p => p.Age).ThenBy(p => p.Name).Last(p => p.Age < 20).Name, + Q.Where(p => p.Age > 10).OrderBy(p => p.Age).ThenBy(p => p.Name).Last(p => p.Age < 20).Name); + Assert.Equal(_people.OrderBy(p => p.Age).ThenBy(p => p.Name).Count(p => p.Age == 7), + Q.OrderBy(p => p.Age).ThenBy(p => p.Name).Count(p => p.Age == 7)); + Assert.Equal(_people.OrderBy(p => p.Age).ThenBy(p => p.Name).LongCount(p => p.Age == 7), + Q.OrderBy(p => p.Age).ThenBy(p => p.Name).LongCount(p => p.Age == 7)); + + // All(p) is not folded and still sees every document. + Assert.True(Q.OrderBy(p => p.Name).All(p => p.Age < 50)); + Assert.False(Q.OrderBy(p => p.Name).All(p => p.Age < 49)); + } + + // Titles run against the vector distance, so ordering by Title picks a different neighbour + // than the order the index returns them in: the 3 nearest to the origin are "j", "i", "h". + private static readonly float[] Origin = [0f, 0f, 0f]; + + private async Task InsertReverseTitledVectorsAsync() + { + for (int i = 0; i < 10; i++) + await _db.VectorItems.InsertAsync(new VectorEntity { Title = ((char)('j' - i)).ToString(), Embedding = [i, 0f, 0f] }); + } + + private IQueryable V => _db.VectorItems.AsQueryable(); + + [Fact] + public async Task ElementTerminals_AfterOrdering_WithVectorSearchPredicate_UseTheFoldedResult() + { + await InsertReverseTitledVectorsAsync(); + + Assert.Equal("h", V.OrderBy(x => x.Title).First(x => x.Embedding.VectorSearch(Origin, 3)).Title); + Assert.Equal(V.Where(x => x.Embedding.VectorSearch(Origin, 3)).OrderBy(x => x.Title).First().Title, + V.OrderBy(x => x.Title).First(x => x.Embedding.VectorSearch(Origin, 3)).Title); + Assert.Equal(V.Where(x => x.Embedding.VectorSearch(Origin, 3)).OrderByDescending(x => x.Title).FirstOrDefault()!.Title, + V.OrderByDescending(x => x.Title).FirstOrDefault(x => x.Embedding.VectorSearch(Origin, 3))!.Title); + Assert.Equal("j", V.OrderBy(x => x.Title).Last(x => x.Embedding.VectorSearch(Origin, 3)).Title); + Assert.Equal(V.Where(x => x.Embedding.VectorSearch(Origin, 3)).OrderBy(x => x.Title).Last().Title, + V.OrderBy(x => x.Title).Last(x => x.Embedding.VectorSearch(Origin, 3)).Title); + Assert.Equal(V.Where(x => x.Embedding.VectorSearch(Origin, 3)).OrderByDescending(x => x.Title).LastOrDefault()!.Title, + V.OrderByDescending(x => x.Title).LastOrDefault(x => x.Embedding.VectorSearch(Origin, 3))!.Title); + Assert.Equal("j", V.OrderBy(x => x.Title).Single(x => x.Embedding.VectorSearch(Origin, 1)).Title); + Assert.Equal("j", V.OrderBy(x => x.Title).SingleOrDefault(x => x.Embedding.VectorSearch(Origin, 1))!.Title); + Assert.Throws(() => V.OrderBy(x => x.Title).Single(x => x.Embedding.VectorSearch(Origin, 3))); + Assert.True(V.OrderBy(x => x.Title).Any(x => x.Embedding.VectorSearch(Origin, 3))); + + // ThenBy is not modelled by the direct pipeline, so these go through the replay as well. + Assert.Equal(V.Where(x => x.Title != null && x.Embedding.VectorSearch(Origin, 3)).OrderBy(x => x.Title.Length).ThenBy(x => x.Title).First().Title, + V.Where(x => x.Title != null).OrderBy(x => x.Title.Length).ThenBy(x => x.Title).First(x => x.Embedding.VectorSearch(Origin, 3)).Title); + Assert.Equal("h", V.OrderBy(x => x.Title.Length).ThenBy(x => x.Title).First(x => x.Embedding.VectorSearch(Origin, 3)).Title); + Assert.Equal(3, V.OrderBy(x => x.Title.Length).ThenBy(x => x.Title).Count(x => x.Embedding.VectorSearch(Origin, 3))); + Assert.Equal(3L, V.OrderBy(x => x.Title.Length).ThenBy(x => x.Title).LongCount(x => x.Embedding.VectorSearch(Origin, 3))); + Assert.True(V.OrderBy(x => x.Title.Length).ThenByDescending(x => x.Title).Any(x => x.Embedding.VectorSearch(Origin, 3))); + } + + [Fact] + public async Task Terminals_WithUnfoldedVectorSearchPredicate_StillThrow() + { + await InsertReverseTitledVectorsAsync(); + + // All(p), and any predicate after an operator that blocks folding, is evaluated in the + // replay, where VectorSearch cannot run. + Assert.Throws(() => V.OrderBy(x => x.Title).All(x => x.Embedding.VectorSearch(Origin, 3))); + Assert.Throws(() => V.OrderBy(x => x.Title).SkipWhile(x => x.Title == "a").First(x => x.Embedding.VectorSearch(Origin, 3))); + Assert.Throws(() => V.OrderBy(x => x.Title).Take(5).Count(x => x.Embedding.VectorSearch(Origin, 3))); + Assert.Throws(() => V.GroupBy(x => x.Title).Count(g => g.First().Embedding.VectorSearch(Origin, 3))); + } + // ── Write path ────────────────────────────────────────────────────────── [Fact]