Make distinct deduplication linear after sorting

Author: xet7

CHANGELOG.md

@@ -6,6 +6,17 @@
 
 ### Fixed 🐛
 
+- **Large `distinct` results no longer perform quadratic BSON
+  deduplication.** The handler previously called linear `Array.Contains` for
+  every value even after SQLite had already collapsed most duplicates, making
+  45,640 unique IDs consume about 15–17 seconds of CPU. Values are now gathered,
+  sorted once as already required by the command, and compacted by comparing
+  adjacent BSON values; the field path is also parsed once per command instead
+  of once per document. The representative 45,640-value compaction benchmark
+  completes in about 5–18 milliseconds. Regression tests cover stable ordering,
+  cross-width numeric equality and duplicate strings, nulls, documents and
+  nested arrays by @xet7. Thanks to xet7.
+
 - **Numeric `$type` checks combined with a negated range no longer decode whole
   SQLite collections.** The unchanged corruption-check selector
   `{sort: {$type: "number"}, $nor: [{sort: {$gte: low, $lte: high}}]}` now

internal/handler/common/distinct.go

@@ -25,6 +25,7 @@ import (
 	"github.com/FerretDB/FerretDB/internal/types"
 	"github.com/FerretDB/FerretDB/internal/util/iterator"
 	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
+	"github.com/FerretDB/FerretDB/internal/util/must"
 )
 
 // DistinctParams contains `distinct` command parameters supported by at least one handler.
@@ -93,7 +94,11 @@ func GetDistinctParams(document *types.Document, l *slog.Logger) (*DistinctParam
 // If the key is found in the document, and the value is an array, each element of the array is added to the result.
 // Otherwise, the value itself is added to the result.
 func FilterDistinctValues(iter types.DocumentsIterator, key string) (*types.Array, error) {
-	distinct := types.MakeArray(0)
+	values := types.MakeArray(0)
+	path, err := types.NewPathFromString(key)
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
 
 	defer iter.Close()
 
@@ -107,11 +112,6 @@ func FilterDistinctValues(iter types.DocumentsIterator, key string) (*types.Arra
 			return nil, lazyerrors.Error(err)
 		}
 
-		path, err := types.NewPathFromString(key)
-		if err != nil {
-			return nil, lazyerrors.Error(err)
-		}
-
 		// distinct using dot notation returns the value by valid array index
 		// or values for the given key in array's document
 		vals, err := commonpath.FindValues(doc, path, &commonpath.FindValuesOpts{
@@ -131,20 +131,36 @@ func FilterDistinctValues(iter types.DocumentsIterator, key string) (*types.Arra
 						return nil, lazyerrors.Error(err)
 					}
 
-					if !distinct.Contains(el) {
-						distinct.Append(el)
-					}
+					values.Append(el)
 				}
 
 			default:
-				if !distinct.Contains(v) {
-					distinct.Append(v)
-				}
+				values.Append(v)
 			}
 		}
 	}
 
-	SortArray(distinct, types.Ascending)
+	return sortAndDeduplicateDistinctValues(values), nil
+
+}
+
+// sortAndDeduplicateDistinctValues replaces the former incremental
+// Array.Contains pass. That pass compared every new value with all preceding
+// values and made a large already-distinct result quadratic. Sorting is already
+// required by the command; equal BSON values become adjacent and need one
+// comparison each during the compaction pass.
+func sortAndDeduplicateDistinctValues(values *types.Array) *types.Array {
+	SortArray(values, types.Ascending)
+	distinct := types.MakeArray(values.Len())
+
+	var previous any
+	for i := 0; i < values.Len(); i++ {
+		value := must.NotFail(values.Get(i))
+		if i == 0 || types.Compare(previous, value) != types.Equal {
+			distinct.Append(value)
+			previous = value
+		}
+	}
 
-	return distinct, nil
+	return distinct
 }

internal/handler/common/distinct_test.go

@@ -0,0 +1,77 @@
+// Copyright 2021 FerretDB 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 common
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/FerretDB/FerretDB/internal/types"
+	"github.com/FerretDB/FerretDB/internal/util/must"
+)
+
+func TestSortAndDeduplicateDistinctValues(t *testing.T) {
+	t.Parallel()
+
+	doc := must.NotFail(types.NewDocument("x", int32(1)))
+	array := must.NotFail(types.NewArray("nested"))
+	values := must.NotFail(types.NewArray(
+		"b", int32(1), doc, "a", types.Null, array,
+		"a", int64(1), doc.DeepCopy(), types.Null, array.DeepCopy(),
+	))
+
+	actual := sortAndDeduplicateDistinctValues(values)
+	require.Equal(t, 6, actual.Len())
+
+	expected := []any{types.Null, int32(1), "a", "b", doc, array}
+	for _, value := range expected {
+		matches := 0
+		for i := 0; i < actual.Len(); i++ {
+			if types.Compare(must.NotFail(actual.Get(i)), value) == types.Equal {
+				matches++
+			}
+		}
+		assert.Equal(t, 1, matches, "value %v must occur exactly once", value)
+	}
+
+	for i := 1; i < actual.Len(); i++ {
+		assert.NotEqual(t, types.Greater, types.CompareOrderForSort(
+			must.NotFail(actual.Get(i-1)), must.NotFail(actual.Get(i)), types.Ascending,
+		))
+	}
+}
+
+func BenchmarkSortAndDeduplicateDistinctValues(b *testing.B) {
+	const values = 45640
+	input := make([]string, values)
+	for i := range input {
+		input[i] = fmt.Sprintf("%024d", values-i)
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		array := types.MakeArray(values)
+		for _, value := range input {
+			array.Append(value)
+		}
+		if got := sortAndDeduplicateDistinctValues(array); got.Len() != values {
+			b.Fatalf("got %d values", got.Len())
+		}
+	}
+}