object/search: Add function merging SearchV2 results

Future use-cases: - merge results from several shard's metabases; - merge results from several SNs. Refs #3058. Signed-off-by: Leonard Lyubich <[email protected]>
nspcc-dev · Feb 7, 2025 · 9ea9553 · 9ea9553
1 parent 400c685
commit 9ea9553
Show file tree

Hide file tree

Showing 2 changed files with 239 additions and 0 deletions.
diff --git a/pkg/core/object/metadata.go b/pkg/core/object/metadata.go
@@ -0,0 +1,76 @@
+package object
+
+import (
+	"bytes"
+	"math/big"
+	"strings"
+
+	"github.com/nspcc-dev/neofs-sdk-go/client"
+)
+
+// MergeSearchResults inserts next set of search result items into pre-allocated
+// buffer with n set items. Returns the number of items added.
+func MergeSearchResults(n uint16, buf, next []client.SearchResultItem) uint16 {
+	switch {
+	case len(buf) == 0:
+		panic("empty buffer")
+	case n > uint16(len(buf)):
+		panic("n overflows buffer len")
+	}
+	if n == 0 {
+		return uint16(copy(buf, next))
+	}
+	withAttr := len(buf[0].Attributes) > 0
+	var iN, iB *big.Int
+	if withAttr {
+		iN, iB = new(big.Int), new(big.Int)
+	}
+	var added uint16
+next:
+	for len(next) > 0 {
+		var isInt bool
+		if withAttr {
+			_, isInt = iN.SetString(next[0].Attributes[0], 10)
+		}
+		var i uint16
+		for i = 0; i < n; i++ { // do not use range: if next[0] is the biggest, i=n is desired, following condition catches
+			if withAttr {
+				if isInt {
+					_, isInt = iB.SetString(buf[i].Attributes[0], 10)
+				}
+				if isInt {
+					if c := iN.Cmp(iB); c < 0 {
+						break
+					} else if c > 0 {
+						continue
+					}
+				} else {
+					if c := strings.Compare(next[0].Attributes[0], buf[i].Attributes[0]); c < 0 {
+						break
+					} else if c > 0 {
+						continue
+					}
+				}
+			}
+			if c := bytes.Compare(next[0].ID[:], buf[i].ID[:]); c == 0 {
+				next = next[1:]
+				continue next
+			} else if c < 0 {
+				break
+			}
+		}
+		if i == uint16(len(buf)) {
+			break
+		}
+		copy(buf[i+1:], buf[i:])
+		buf[i] = next[0]
+		added++
+		next, buf = next[1:], buf[i+1:]
+		if i < n {
+			n -= i + 1
+		} else {
+			n = 0
+		}
+	}
+	return added
+}
diff --git a/pkg/core/object/metadata_test.go b/pkg/core/object/metadata_test.go
@@ -0,0 +1,163 @@
+package object_test
+
+import (
+	"bytes"
+	"slices"
+	"strconv"
+	"strings"
+	"testing"
+
+	. "github.com/nspcc-dev/neofs-node/pkg/core/object"
+	"github.com/nspcc-dev/neofs-sdk-go/client"
+	oidtest "github.com/nspcc-dev/neofs-sdk-go/object/id/test"
+	"github.com/stretchr/testify/require"
+)
+
+func searchResultFromIDs(n int) []client.SearchResultItem {
+	ids := oidtest.IDs(n)
+	s := make([]client.SearchResultItem, len(ids))
+	for i := range ids {
+		s[i].ID = ids[i]
+	}
+	sortSearchResult(s)
+	return s
+}
+
+func sortSearchResult(s []client.SearchResultItem) {
+	slices.SortFunc(s, func(a, b client.SearchResultItem) int {
+		if len(a.Attributes) > 0 {
+			if c := strings.Compare(a.Attributes[0], b.Attributes[0]); c != 0 {
+				return c
+			}
+		}
+		return bytes.Compare(a.ID[:], b.ID[:])
+	})
+}
+
+func TestMergeSearchResults(t *testing.T) {
+	anyValidBuf := []client.SearchResultItem{{}}
+	t.Run("empty buffer", func(t *testing.T) {
+		require.PanicsWithValue(t, "empty buffer", func() { MergeSearchResults(0, nil, nil) })
+		require.PanicsWithValue(t, "empty buffer", func() { MergeSearchResults(0, []client.SearchResultItem{}, nil) })
+	})
+	t.Run("collected num overflow", func(t *testing.T) {
+		require.PanicsWithValue(t, "n overflows buffer len", func() { MergeSearchResults(2, anyValidBuf, nil) })
+	})
+	t.Run("no attributes in next item", func(t *testing.T) {
+		require.Panics(t, func() {
+			MergeSearchResults(1, []client.SearchResultItem{
+				{ID: oidtest.ID(), Attributes: []string{"any"}},
+			}, []client.SearchResultItem{
+				{ID: oidtest.ID(), Attributes: nil},
+			})
+		})
+	})
+	t.Run("add nothing", func(t *testing.T) {
+		require.Zero(t, MergeSearchResults(1, anyValidBuf, nil))
+		require.Zero(t, MergeSearchResults(1, anyValidBuf, []client.SearchResultItem{}))
+	})
+	test := func(t testing.TB, modify func([]client.SearchResultItem)) {
+		b := make([]client.SearchResultItem, 10)
+		s := searchResultFromIDs(len(b) + 2)
+		modify(s)
+		// first add from the middle
+		n := MergeSearchResults(0, b, s[4:7])
+		require.EqualValues(t, 3, n)
+		require.Equal(t, s[4:7], b[:n])
+		// add some of them again + append 1 item
+		n += MergeSearchResults(n, b, s[5:8])
+		require.EqualValues(t, 4, n)
+		require.Equal(t, s[4:8], b[:n])
+		// unshift 2 smallest items
+		n += MergeSearchResults(n, b, s[:2])
+		require.EqualValues(t, 6, n)
+		require.Equal(t, slices.Concat(s[:2], s[4:8]), b[:n])
+		// append 2 biggest items
+		n += MergeSearchResults(n, b, s[len(s)-2:])
+		require.EqualValues(t, 8, n)
+		require.Equal(t, slices.Concat(s[:2], s[4:8], s[10:]), b[:n])
+		// unshift one more
+		n += MergeSearchResults(n, b, s[3:5])
+		require.EqualValues(t, 9, n)
+		require.Equal(t, slices.Concat(s[:2], s[3:8], s[10:]), b[:n])
+		// unshift one more
+		n += MergeSearchResults(n, b, s[2:3])
+		require.EqualValues(t, 10, n)
+		require.Equal(t, slices.Concat(s[:8], s[10:]), b[:n])
+		// insert one item so the biggest one goes out
+		require.EqualValues(t, 1, MergeSearchResults(n, b, s[9:10]))
+		require.Equal(t, slices.Concat(s[:8], s[9:11]), b[:n])
+		// and again
+		require.EqualValues(t, 1, MergeSearchResults(n, b, s[5:9]))
+		require.Equal(t, s[:10], b[:n])
+		// add the biggest ones again
+		require.Zero(t, MergeSearchResults(n, b, s[10:]))
+		require.Equal(t, s[:10], b[:n])
+		// add same items again
+		for i := range 11 {
+			require.Zero(t, MergeSearchResults(n, b, []client.SearchResultItem{s[i]}))
+			require.Equal(t, s[:10], b[:n])
+			require.Zero(t, MergeSearchResults(n, b, []client.SearchResultItem{s[i], s[i]}))
+			require.Equal(t, s[:10], b[:n])
+			require.Zero(t, MergeSearchResults(n, b, s[:i]))
+			require.Equal(t, s[:10], b[:n])
+			require.Zero(t, MergeSearchResults(n, b, s[11-i:]))
+			require.Equal(t, s[:10], b[:n])
+		}
+	}
+	t.Run("no attributes", func(t *testing.T) {
+		test(t, func([]client.SearchResultItem) {})
+	})
+	test(t, func(s []client.SearchResultItem) {
+		for i := range s {
+			n := i
+			if i >= 3 && i <= 5 {
+				n = 3
+			}
+			s[i].Attributes = []string{"val1_" + strconv.Itoa(n), "val2_" + strconv.Itoa(i)}
+		}
+		sortSearchResult(s)
+	})
+	t.Run("integers", func(t *testing.T) {
+		b := make([]client.SearchResultItem, 5)
+		s := searchResultFromIDs(7)
+		slices.Reverse(s)
+		for i, a := range [7]string{
+			"-111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111",
+			"-18446744073709551615",
+			"-1",
+			"0",
+			"1",
+			"18446744073709551615",
+			"111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111",
+		} {
+			s[i].Attributes = []string{a}
+		}
+		n := MergeSearchResults(0, b, slices.Concat([]client.SearchResultItem{s[1]}, s[4:]))
+		require.EqualValues(t, 4, n)
+		require.Equal(t, slices.Concat([]client.SearchResultItem{s[1]}, s[4:]), b[:n])
+		// insert zero
+		n += MergeSearchResults(n, b, []client.SearchResultItem{s[3]})
+		require.EqualValues(t, 5, n)
+		require.Equal(t, slices.Concat([]client.SearchResultItem{s[1]}, s[3:]), b[:n])
+		// unshift lowest
+		require.EqualValues(t, 1, MergeSearchResults(n, b, []client.SearchResultItem{s[0]}))
+		require.Equal(t, slices.Concat(s[:2], s[3:6]), b[:n])
+		// insert all
+		require.EqualValues(t, 1, MergeSearchResults(n, b, s))
+		require.Equal(t, s[:len(b)], b)
+		require.EqualValues(t, 0, MergeSearchResults(5, b, s))
+		// insert smaller non-int
+		next := searchResultFromIDs(2)
+		next[0].Attributes = []string{"!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"}
+		next[1].Attributes = []string{"++++++++++++++++++++++++++"}
+		require.EqualValues(t, 2, MergeSearchResults(5, b, next))
+		require.Equal(t, slices.Concat(next, s[:3]), b)
+		// insert bigger non-int
+		next2 := searchResultFromIDs(2)
+		next2[0].Attributes = []string{"2a"}
+		next2[1].Attributes = []string{"3a"}
+		require.EqualValues(t, 0, MergeSearchResults(5, b, next2))
+		require.Equal(t, slices.Concat(next, s[:3]), b)
+	})
+}