diff --git a/internal/backends/dolt/collated_index_filter_test.go b/internal/backends/dolt/collated_index_filter_test.go new file mode 100644 index 00000000..e6fabf2e --- /dev/null +++ b/internal/backends/dolt/collated_index_filter_test.go @@ -0,0 +1,359 @@ +// Copyright 2026 Dolthub, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dolt + +import ( + "context" + "fmt" + "reflect" + "sort" + "testing" + + "github.com/dolthub/dumbodb/internal/backends" + "github.com/dolthub/dumbodb/internal/collation" + "github.com/dolthub/dumbodb/internal/handler/common" + "github.com/dolthub/dumbodb/internal/types" + "github.com/dolthub/dumbodb/internal/util/must" +) + +// TestCollatedIndexOperatorFilters is the issue #130 reproduction: operator +// filters on a collated index field panicked with "unexpected type +// index.PreEncoded". +func TestCollatedIndexOperatorFilters(t *testing.T) { + t.Parallel() + + ctx := context.Background() + b := newTestBackend(t) + db, err := b.Database("collop") + if err != nil { + t.Fatalf("Database: %v", err) + } + coll, err := db.Collection("c") + if err != nil { + t.Fatalf("Collection: %v", err) + } + + const n = 100 + docs := make([]*types.Document, 0, n) + for i := 0; i < n; i++ { + docs = append(docs, must.NotFail(types.NewDocument( + "_id", int32(i), + "s", fmt.Sprintf("s%03d", i), + "i", int32(i), + ))) + } + if _, err := coll.InsertAll(ctx, &backends.InsertAllParams{Docs: docs}); err != nil { + t.Fatalf("InsertAll: %v", err) + } + + frCA := mustDoc(t, "locale", "fr_CA") + if _, err := coll.CreateIndexes(ctx, &backends.CreateIndexesParams{ + Indexes: []backends.IndexInfo{{ + Name: "s_idx", + Key: []backends.IndexKeyPair{{Field: "s"}}, + Collation: frCA, + }}, + }); err != nil { + t.Fatalf("CreateIndexes: %v", err) + } + + // lookup returns the backend's candidate _ids; the handler re-filters them. + lookup := func(t *testing.T, filter *types.Document) []int32 { + t.Helper() + got := drainQuery(t, ctx, coll, &backends.QueryParams{ + Filter: filter, + Collated: true, + Collation: frCA, + }) + ids := idsInDocs(t, got) + sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) + return ids + } + + t.Run("gt", func(t *testing.T) { + op := must.NotFail(types.NewDocument("$gt", "s090")) + ids := lookup(t, must.NotFail(types.NewDocument("s", op))) + want := make([]int32, 0, 9) + for i := int32(91); i < n; i++ { + want = append(want, i) + } + if len(ids) != 9 || !reflect.DeepEqual(ids, want) { + t.Errorf("$gt s090 = %v, want [91..99] (9 docs)", ids) + } + }) + + t.Run("eq", func(t *testing.T) { + op := must.NotFail(types.NewDocument("$eq", "s050")) + ids := lookup(t, must.NotFail(types.NewDocument("s", op))) + if len(ids) != 1 || !reflect.DeepEqual(ids, []int32{50}) { + t.Errorf("$eq s050 = %v, want [50]", ids) + } + }) + + t.Run("bounded_range", func(t *testing.T) { + op := must.NotFail(types.NewDocument("$gte", "s010", "$lt", "s020")) + ids := lookup(t, must.NotFail(types.NewDocument("s", op))) + want := make([]int32, 0, 10) + for i := int32(10); i < 20; i++ { + want = append(want, i) + } + if len(ids) != 10 || !reflect.DeepEqual(ids, want) { + t.Errorf("$gte s010, $lt s020 = %v, want [10..19] (10 docs)", ids) + } + }) + + // Non-string operands are not collated. On s they bracket below every + // string key; on the binary-indexed i the collated query must not use + // i_idx, so it scans the whole collection. + t.Run("nonstring_operator_values", func(t *testing.T) { + op := must.NotFail(types.NewDocument("$gt", int32(90))) + sIDs := lookup(t, must.NotFail(types.NewDocument("s", op))) + if len(sIDs) != 0 { + t.Errorf("$gt 90 candidates on collated s = %v, want empty (0 docs)", sIDs) + } + + if _, err := coll.CreateIndexes(ctx, &backends.CreateIndexesParams{ + Indexes: []backends.IndexInfo{{ + Name: "i_idx", + Key: []backends.IndexKeyPair{{Field: "i"}}, + }}, + }); err != nil { + t.Fatalf("CreateIndexes (binary): %v", err) + } + + opI := must.NotFail(types.NewDocument("$gte", int32(10), "$lt", int32(20))) + iIDs := lookup(t, must.NotFail(types.NewDocument("i", opI))) + if len(iIDs) != n { + t.Errorf("binary-field numeric range candidates = %d, want %d (whole-collection scan candidates; the collated query must not use the binary index)", len(iIDs), n) + } + }) + + // A numeric and a string lower bound merge by BSON type order in either + // insertion order, leaving the collated string bound. + t.Run("mixed_type_lower_bounds", func(t *testing.T) { + op := must.NotFail(types.NewDocument("$gte", int32(5), "$gt", "s090")) + ids := lookup(t, must.NotFail(types.NewDocument("s", op))) + want := make([]int32, 0, 9) + for i := int32(91); i < n; i++ { + want = append(want, i) + } + if len(ids) != 9 || !reflect.DeepEqual(ids, want) { + t.Errorf("$gte 5, $gt s090 candidates = %v, want [91..99] (9 docs)", ids) + } + }) + + t.Run("mixed_type_lower_bounds_reversed", func(t *testing.T) { + op := must.NotFail(types.NewDocument("$gt", "s090", "$gte", int32(5))) + ids := lookup(t, must.NotFail(types.NewDocument("s", op))) + want := make([]int32, 0, 9) + for i := int32(91); i < n; i++ { + want = append(want, i) + } + if len(ids) != 9 || !reflect.DeepEqual(ids, want) { + t.Errorf("$gt s090, $gte 5 candidates = %v, want [91..99] (9 docs)", ids) + } + }) +} + +// TestCollatedIndexOperatorSeeksFewerNodes checks, at the seek level (not just +// "no panic"), that an operator filter on a collated index is served by the +// sort-key index: it must read fewer prolly-tree nodes than a full scan. +func TestCollatedIndexOperatorSeeksFewerNodes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + b := newTestBackend(t) + db, err := b.Database("collopseek") + if err != nil { + t.Fatalf("Database: %v", err) + } + coll, err := db.Collection("c") + if err != nil { + t.Fatalf("Collection: %v", err) + } + + const n = 5000 + docs := make([]*types.Document, 0, n) + for i := 0; i < n; i++ { + docs = append(docs, must.NotFail(types.NewDocument( + "_id", int32(i), + "s", fmt.Sprintf("s%04d", i), + ))) + } + if _, err := coll.InsertAll(ctx, &backends.InsertAllParams{Docs: docs}); err != nil { + t.Fatalf("InsertAll: %v", err) + } + + frCA := mustDoc(t, "locale", "fr_CA") + if _, err := coll.CreateIndexes(ctx, &backends.CreateIndexesParams{ + Indexes: []backends.IndexInfo{{ + Name: "s_idx", + Key: []backends.IndexKeyPair{{Field: "s"}}, + Collation: frCA, + }}, + }); err != nil { + t.Fatalf("CreateIndexes: %v", err) + } + + seekCtx, seek := WithSeekCounter(ctx) + seekDocs := drainQuery(t, seekCtx, coll, &backends.QueryParams{ + Filter: must.NotFail(types.NewDocument("s", + must.NotFail(types.NewDocument("$gte", "s0100", "$lt", "s0110")))), + Collated: true, + Collation: frCA, + }) + if len(seekDocs) != 10 { + t.Errorf("collated range candidates = %d, want 10", len(seekDocs)) + } + + scanCtx, scan := WithSeekCounter(ctx) + allDocs := drainQuery(t, scanCtx, coll, &backends.QueryParams{ + Filter: must.NotFail(types.NewDocument()), + }) + if len(allDocs) != n { + t.Fatalf("full scan returned %d docs, want %d", len(allDocs), n) + } + + t.Logf("collated range nodes=%d; full scan nodes=%d", seek.Nodes(), scan.Nodes()) + + if seek.Nodes() == 0 { + t.Fatal("instrumentation dead: collated seek counted 0 node reads") + } + if scan.Nodes() == 0 { + t.Fatal("instrumentation dead: full scan counted 0 node reads") + } + if seek.Nodes() >= scan.Nodes() { + t.Fatalf("collated operator seek not served by the index: seek nodes=%d >= scan nodes=%d", + seek.Nodes(), scan.Nodes()) + } +} + +// TestCollatedIndexOperatorFiltersCollationOrder uses case and accent variants, +// whose collated order differs from their binary order, so a bound built from +// the wrong key space drops or admits documents. Numeric fillers keep string +// ranges selective enough that the planner does not fall back to a scan. +func TestCollatedIndexOperatorFiltersCollationOrder(t *testing.T) { + t.Parallel() + + ctx := context.Background() + vals := []any{ + "a", "A", "ab", "aB", "Ab", "b", "B", "c", "C", "cote", "cot\u00e9", "c\u00f4te", "c\u00f4t\u00e9", + "e", "E", "\u00e9", "z", "Z", "", int32(5), int32(50), true, + } + const fillers = 200 + op := func(kv ...any) *types.Document { + return must.NotFail(types.NewDocument("s", must.NotFail(types.NewDocument(kv...)))) + } + arr := func(v ...any) *types.Array { return must.NotFail(types.NewArray(v...)) } + + // exact: single-type string bounds, where the index range is the match set. + type filterCase struct { + name string + filter *types.Document + exact bool + } + cases := []filterCase{ + {"eq_B", op("$eq", "B"), true}, + {"gt_b", op("$gt", "b"), true}, + {"gte_B", op("$gte", "B"), true}, + {"lt_c", op("$lt", "c"), true}, + {"lte_C", op("$lte", "C"), true}, + {"range_b_d", op("$gte", "b", "$lt", "d"), true}, + {"range_cote", op("$gt", "cote", "$lte", "c\u00f4t\u00e9"), true}, + {"range_E_f", op("$gte", "E", "$lt", "f"), true}, + {"in_B_z", op("$in", arr("B", "z")), false}, + {"in_mixed", op("$in", arr("A", int32(5))), false}, + {"ne_B", op("$ne", "B"), false}, + {"not_gt_b", op("$not", must.NotFail(types.NewDocument("$gt", "b"))), false}, + {"mixed_lower", op("$gte", int32(5), "$gt", "b"), false}, + {"mixed_upper", op("$lt", int32(60), "$lte", "c"), false}, + {"eq_and_gte", op("$eq", "b", "$gte", "B"), false}, + } + + for _, spec := range [][]any{ + {"locale", "en", "strength", int32(1)}, + {"locale", "en", "strength", int32(2)}, + {"locale", "fr_CA"}, + } { + collDoc := must.NotFail(types.NewDocument(spec...)) + t.Run(fmt.Sprint(spec), func(t *testing.T) { + t.Parallel() + + b := newTestBackend(t) + db := must.NotFail(b.Database("collorder")) + coll := must.NotFail(db.Collection("c")) + docs := make([]*types.Document, 0, len(vals)+fillers) + for i, v := range vals { + docs = append(docs, must.NotFail(types.NewDocument("_id", int32(i), "s", v))) + } + for i := len(vals); i < len(vals)+fillers; i++ { + docs = append(docs, must.NotFail(types.NewDocument("_id", int32(i), "s", int32(1000+i)))) + } + if _, err := coll.InsertAll(ctx, &backends.InsertAllParams{Docs: docs}); err != nil { + t.Fatalf("InsertAll: %v", err) + } + if _, err := coll.CreateIndexes(ctx, &backends.CreateIndexesParams{ + Indexes: []backends.IndexInfo{{ + Name: "s_idx", + Key: []backends.IndexKeyPair{{Field: "s"}}, + Collation: collDoc, + }}, + }); err != nil { + t.Fatalf("CreateIndexes: %v", err) + } + cmp := collation.Parse(collDoc).Comparator() + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var want []int32 + for _, d := range docs { + if must.NotFail(common.FilterDocumentColl(d, tc.filter, cmp)) { + want = append(want, must.NotFail(d.Get("_id")).(int32)) + } + } + got := idsInDocs(t, drainQuery(t, ctx, coll, &backends.QueryParams{ + Filter: tc.filter, + Collated: true, + Collation: collDoc, + })) + sort.Slice(got, func(i, j int) bool { return got[i] < got[j] }) + + candidates := make(map[int32]bool, len(got)) + for _, id := range got { + candidates[id] = true + } + for _, id := range want { + if !candidates[id] { + t.Errorf("candidates %v miss match _id %d", got, id) + } + } + if !tc.exact { + return + } + if !reflect.DeepEqual(got, want) && !(len(got) == 0 && len(want) == 0) { + t.Errorf("candidates = %v, want %v", got, want) + } + res, err := coll.Count(ctx, &backends.CountParams{Filter: tc.filter, Collation: collDoc}) + if err != nil { + t.Fatalf("Count: %v", err) + } + if res.Filtered && res.Count != int64(len(want)) { + t.Errorf("indexed count = %d, want %d", res.Count, len(want)) + } + }) + } + }) + } +} diff --git a/internal/backends/dolt/collection.go b/internal/backends/dolt/collection.go index e310ac2c..7ed2b125 100644 --- a/internal/backends/dolt/collection.go +++ b/internal/backends/dolt/collection.go @@ -957,7 +957,10 @@ func multikeyScanBound(v any) any { lower := types.MakeDocument(2) for _, op := range opDoc.Keys() { if op == "$gt" || op == "$gte" { - lower.Set(op, must.NotFail(opDoc.Get(op))) + // AppendDecoded: the operand may be a collated sort key + // (idxpkg.PreEncoded), an internal planner value. + v, _ := opDoc.Get(op) + lower.AppendDecoded(op, v) } } return lower @@ -1343,7 +1346,20 @@ func compoundIndexBounds(idx backends.IndexInfo, filter *types.Document, queryCm // compareScalars returns -1 / 0 / 1 using types.Compare. Used to merge // multiple bounds on the same side ($gte:1, $gte:5 -> keep the tighter 5). +// Collated sort keys (idxpkg.PreEncoded) compare bytewise; against a +// non-string operand ({$gte: 5, $gt: "b"}) a key stands in as "" so BSON +// type order decides. func compareScalars(a, b any) int { + pa, aKey := a.(idxpkg.PreEncoded) + pb, bKey := b.(idxpkg.PreEncoded) + switch { + case aKey && bKey: + return bytes.Compare(pa, pb) + case aKey: + a = "" + case bKey: + b = "" + } switch types.Compare(a, b) { case types.Less: return -1 @@ -2310,10 +2326,10 @@ func collateFilterValue(v any, cmp *collation.Comparator) any { for _, k := range t.Keys() { ov, _ := t.Get(k) if str, ok := ov.(string); ok { - out.Set(k, idxpkg.PreEncoded(cmp.Key(str))) - } else { - out.Set(k, ov) + ov = idxpkg.PreEncoded(cmp.Key(str)) } + // AppendDecoded: Set rejects idxpkg.PreEncoded. + out.AppendDecoded(k, ov) } return out default: diff --git a/internal/index/keystring.go b/internal/index/keystring.go index 2503d8ff..867fbcc3 100644 --- a/internal/index/keystring.go +++ b/internal/index/keystring.go @@ -489,7 +489,7 @@ func BracketRange(v any) (start, stop []byte, ok bool) { case int32, int64, float64: // [-Inf .. +Inf], excluding NaN below it. return []byte{ctypeNegInf}, []byte{ctypePosInf + 1}, true - case string: + case string, PreEncoded: return []byte{ctypeString}, []byte{ctypeString + 1}, true case time.Time: return []byte{ctypeDate}, []byte{ctypeDate + 1}, true