fix: collection filtered pagination (backport #23002) (#23139)

Co-authored-by: Akhil Kumar P <36399231+akhilkumarpilli@users.noreply.github.com>
Co-authored-by: akhilkumarpilli <akhilkumar7947@gmail.com>
Co-authored-by: Julien Robert <julien@rbrt.fr>

types/query/collections_pagination.go

@@ -266,53 +266,59 @@ func collFilteredPaginateByKey[K, V any, C Collection[K, V], T any](
 	defer iterator.Close()
 
 	var (
-		count   uint64
-		nextKey []byte
+		count       uint64
+		nextKey     []byte
+		transformed T
 	)
 
 	for ; iterator.Valid(); iterator.Next() {
-		// if we reached the specified limit
-		// then we get the next key, and we exit the iteration.
-		if count == limit {
-			concreteKey, err := iterator.Key()
-			if err != nil {
-				return nil, nil, err
-			}
-
-			nextKey, err = encodeCollKey[K, V](coll, concreteKey)
-			if err != nil {
-				return nil, nil, err
-			}
-			break
-		}
-
 		kv, err := iterator.KeyValue()
 		if err != nil {
 			return nil, nil, err
 		}
+
+		include := false
 		// if no predicate is specified then we just append the result
 		if predicateFunc == nil {
-			transformed, err := transformFunc(kv.Key, kv.Value)
+			transformed, err = transformFunc(kv.Key, kv.Value)
 			if err != nil {
 				return nil, nil, err
 			}
-			results = append(results, transformed)
+			include = true
 			// if predicate is applied we execute the predicate function
 			// and append only if predicateFunc yields true.
 		} else {
-			include, err := predicateFunc(kv.Key, kv.Value)
+			include, err = predicateFunc(kv.Key, kv.Value)
 			if err != nil {
 				return nil, nil, err
 			}
 			if include {
-				transformed, err := transformFunc(kv.Key, kv.Value)
+				transformed, err = transformFunc(kv.Key, kv.Value)
 				if err != nil {
 					return nil, nil, err
 				}
-				results = append(results, transformed)
 			}
 		}
-		count++
+
+		if include {
+			// if we reached the specified limit
+			// then we get the next key, and we exit the iteration.
+			if count == limit {
+				concreteKey, err := iterator.Key()
+				if err != nil {
+					return nil, nil, err
+				}
+
+				nextKey, err = encodeCollKey[K, V](coll, concreteKey)
+				if err != nil {
+					return nil, nil, err
+				}
+				break
+			}
+
+			results = append(results, transformed)
+			count++
+		}
 	}
 
 	return results, &PageResponse{

types/query/collections_pagination_test.go

@@ -127,7 +127,7 @@ func TestCollectionPagination(t *testing.T) {
 				Limit: 3,
 			},
 			expResp: &PageResponse{
-				NextKey: encodeKey(5),
+				NextKey: encodeKey(8),
 			},
 			filter: func(key, value uint64) (bool, error) {
 				return key%2 == 0, nil
@@ -135,6 +135,21 @@ func TestCollectionPagination(t *testing.T) {
 			expResults: []collections.KeyValue[uint64, uint64]{
 				{Key: 2, Value: 2},
 				{Key: 4, Value: 4},
+				{Key: 6, Value: 6},
+			},
+		},
+		"filtered with key and empty next key in response": {
+			req: &PageRequest{
+				Key: encodeKey(295),
+			},
+			expResp: &PageResponse{
+				NextKey: nil,
+			},
+			filter: func(key, value uint64) (bool, error) {
+				return key%5 == 0, nil
+			},
+			expResults: []collections.KeyValue[uint64, uint64]{
+				{Key: 295, Value: 295},
 			},
 		},
 	}