diff --git a/src/BLite.Core/Collections/DocumentCollection.cs b/src/BLite.Core/Collections/DocumentCollection.cs index 5893419..605ed62 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/Indexing/BTreeIndex.cs b/src/BLite.Core/Indexing/BTreeIndex.cs index db20278..d7ac1e2 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 3d55687..43866a2 100644 --- a/src/BLite.Core/Query/BTreeExpressionVisitor.cs +++ b/src/BLite.Core/Query/BTreeExpressionVisitor.cs @@ -44,6 +44,15 @@ protected override Expression VisitMethodCall(MethodCallExpression node) case "LongCount": VisitCount(node); return node; + case "First": + case "FirstOrDefault": + case "Single": + case "SingleOrDefault": + case "Last": + case "LastOrDefault": + case "Any": + VisitElementTerminal(node); + return node; default: // GroupBy, Join, etc. — operators the direct pipeline cannot model. _model.HasComplexOperators = true; @@ -60,8 +69,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 +156,36 @@ private void VisitAggregate(MethodCallExpression node) } } + /// + /// 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) + { + 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 } + && CanFoldPredicate(node.Arguments[0], predicate)) + { + AddPredicate(predicate); + _model.ElementTerminal = node.Method.Name; + _model.FoldedTerminal = node; + } + 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,10 +199,57 @@ private void VisitCount(MethodCallExpression node) // Count() with no predicate — fast O(log n) + k BTree scan. _model.IsCountOnly = true; } + 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(predicate); + _model.IsCountOnly = true; + _model.FoldedTerminal = node; + } else { - // Count(predicate) — 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/Query/BTreeQueryProvider.cs b/src/BLite.Core/Query/BTreeQueryProvider.cs index 32965ad..518e75c 100644 --- a/src/BLite.Core/Query/BTreeQueryProvider.cs +++ b/src/BLite.Core/Query/BTreeQueryProvider.cs @@ -214,6 +214,35 @@ 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, ...), and Last / LastOrDefault, which need every match, + // go through the general fallback. + if (model.ElementTerminal is { } terminal) + { + 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"; + 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 @@ -314,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); @@ -675,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 e998fad..ed06955 100644 --- a/src/BLite.Core/Query/QueryModel.cs +++ b/src/BLite.Core/Query/QueryModel.cs @@ -34,4 +34,17 @@ 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, 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/src/BLite.Core/Transactions/Transaction.cs b/src/BLite.Core/Transactions/Transaction.cs index 4dfd990..41d6b72 100644 --- a/src/BLite.Core/Transactions/Transaction.cs +++ b/src/BLite.Core/Transactions/Transaction.cs @@ -50,18 +50,18 @@ 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) { 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/NearWithinCorrectnessTests.cs b/tests/BLite.Tests/NearWithinCorrectnessTests.cs index 9fc5ae6..6721a4e 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 ─────────────────────────────────────────────────── diff --git a/tests/BLite.Tests/QueryAndWritePathTests.cs b/tests/BLite.Tests/QueryAndWritePathTests.cs new file mode 100644 index 0000000..8be51ad --- /dev/null +++ b/tests/BLite.Tests/QueryAndWritePathTests.cs @@ -0,0 +1,339 @@ +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); + + 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() + { + 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)); + } + + [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); + Assert.Null(afterSkipWhile.FoldedTerminal); + + 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))); + } + + [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] + 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(); + } +}