github.com/nutsdb/nutsdb@v1.0.4/tx_zset_test.go (about) 1 // Copyright 2023 The nutsdb Authors. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package nutsdb 16 17 import ( 18 "fmt" 19 "testing" 20 21 "github.com/stretchr/testify/assert" 22 "github.com/stretchr/testify/require" 23 ) 24 25 var tx *Tx 26 27 func TestTx_ZCheck(t *testing.T) { 28 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 29 err := db.View(func(tx *Tx) error { 30 require.Equal(t, ErrBucketNotExist, tx.ZCheck("fake bucket")) 31 return nil 32 }) 33 require.NoError(t, err) 34 }) 35 } 36 37 func TestTx_ZAdd(t *testing.T) { 38 39 bucket := "bucket" 40 41 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 42 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 43 for i := 0; i < 10; i++ { 44 txZAdd(t, db, bucket, GetTestBytes(0), GetTestBytes(i), float64(i), nil, nil) 45 } 46 47 txZCard(t, db, bucket, GetTestBytes(0), 10, nil) 48 }) 49 } 50 51 func TestTx_ZScore(t *testing.T) { 52 bucket := "bucket" 53 54 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 55 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 56 for i := 0; i < 10; i++ { 57 txZAdd(t, db, bucket, GetTestBytes(0), GetTestBytes(i), float64(i), nil, nil) 58 } 59 60 for i := 0; i < 10; i++ { 61 txZScore(t, db, bucket, GetTestBytes(0), GetTestBytes(i), float64(i), nil) 62 } 63 64 txZScore(t, db, bucket, GetTestBytes(0), GetTestBytes(10), float64(10), ErrSortedSetMemberNotExist) 65 txZScore(t, db, bucket, GetTestBytes(1), GetTestBytes(0), float64(0), ErrSortedSetNotFound) 66 67 // update the score of member 68 txZAdd(t, db, bucket, GetTestBytes(0), GetTestBytes(5), float64(999), nil, nil) 69 txZScore(t, db, bucket, GetTestBytes(0), GetTestBytes(5), 999, nil) 70 }) 71 } 72 73 func TestTx_ZRem(t *testing.T) { 74 bucket := "bucket" 75 76 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 77 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 78 79 for i := 0; i < 10; i++ { 80 txZAdd(t, db, bucket, GetTestBytes(0), GetTestBytes(i), float64(i), nil, nil) 81 } 82 83 txZScore(t, db, bucket, GetTestBytes(0), GetTestBytes(3), float64(3), nil) 84 85 // normal remove 86 txZRem(t, db, bucket, GetTestBytes(0), GetTestBytes(3), nil) 87 txZScore(t, db, bucket, GetTestBytes(0), GetTestBytes(3), float64(3), ErrSortedSetMemberNotExist) 88 89 txZCard(t, db, bucket, GetTestBytes(0), 9, nil) 90 91 // remove a fake member 92 txZRem(t, db, bucket, GetTestBytes(0), GetTestBytes(10), ErrSortedSetMemberNotExist) 93 94 // remove from a fake zset 95 txZRem(t, db, bucket, GetTestBytes(1), GetTestBytes(0), ErrSortedSetNotFound) 96 }) 97 } 98 99 func TestTx_ZMembers(t *testing.T) { 100 101 bucket := "bucket" 102 key := GetTestBytes(0) 103 104 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 105 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 106 107 for i := 0; i < 10; i++ { 108 txZAdd(t, db, bucket, key, GetTestBytes(i), float64(i), nil, nil) 109 } 110 111 err := db.View(func(tx *Tx) error { 112 members, err := tx.ZMembers(bucket, key) 113 require.NoError(t, err) 114 115 require.Len(t, members, 10) 116 117 for member := range members { 118 require.Equal(t, GetTestBytes(int(member.Score)), member.Value) 119 } 120 121 return nil 122 }) 123 require.NoError(t, err) 124 }) 125 } 126 127 func TestTx_ZCount(t *testing.T) { 128 129 bucket := "bucket" 130 key := GetTestBytes(0) 131 132 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 133 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 134 135 for i := 0; i < 30; i++ { 136 txZAdd(t, db, bucket, key, GetRandomBytes(24), float64(i), nil, nil) 137 } 138 139 err := db.View(func(tx *Tx) error { 140 141 count, err := tx.ZCount(bucket, key, 10, 20, nil) 142 require.NoError(t, err) 143 require.Equal(t, 11, count) 144 145 return nil 146 }) 147 require.NoError(t, err) 148 }) 149 } 150 151 func TestTx_ZPop(t *testing.T) { 152 bucket := "bucket" 153 key := GetTestBytes(0) 154 155 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 156 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 157 158 txZPop(t, db, bucket, key, true, nil, 0, ErrSortedSetNotFound) 159 txZPop(t, db, bucket, key, false, nil, 0, ErrSortedSetNotFound) 160 161 txZAdd(t, db, bucket, key, GetTestBytes(0), float64(0), nil, nil) 162 txZRem(t, db, bucket, key, GetTestBytes(0), nil) 163 164 txZPop(t, db, bucket, key, true, nil, 0, ErrSortedSetIsEmpty) 165 txZPop(t, db, bucket, key, false, nil, 0, ErrSortedSetIsEmpty) 166 167 for i := 0; i < 30; i++ { 168 txZAdd(t, db, bucket, key, GetTestBytes(i), float64(i), nil, nil) 169 } 170 171 txZPop(t, db, bucket, key, true, GetTestBytes(29), float64(29), nil) 172 txZPop(t, db, bucket, key, false, GetTestBytes(0), 0, nil) 173 174 txZPop(t, db, bucket, key, true, GetTestBytes(28), float64(28), nil) 175 txZPop(t, db, bucket, key, false, GetTestBytes(1), 1, nil) 176 }) 177 } 178 179 func TestTx_ZRangeByScore(t *testing.T) { 180 bucket := "bucket" 181 key := GetTestBytes(0) 182 183 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 184 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 185 err := db.View((func(tx *Tx) error { 186 _, err := tx.ZRangeByScore(bucket, key, 1, 10, nil) 187 require.Error(t, err) 188 return nil 189 })) 190 require.NoError(t, err) 191 for i := 0; i < 10; i++ { 192 txZAdd(t, db, bucket, key, GetTestBytes(i), float64(i), nil, nil) 193 } 194 195 err = db.View(func(tx *Tx) error { 196 members, err := tx.ZRangeByScore(bucket, key, 0, 11, nil) 197 require.NoError(t, err) 198 require.Len(t, members, 10) 199 return nil 200 }) 201 require.NoError(t, err) 202 err = db.View(func(tx *Tx) error { 203 members, err := tx.ZRangeByScore(bucket, key, -1, 2, nil) 204 require.NoError(t, err) 205 require.Len(t, members, 3) 206 return nil 207 }) 208 require.NoError(t, err) 209 err = db.View(func(tx *Tx) error { 210 members, err := tx.ZRangeByScore(bucket, key, 8, 12, nil) 211 require.NoError(t, err) 212 require.Len(t, members, 2) 213 return nil 214 }) 215 require.NoError(t, err) 216 err = db.View(func(tx *Tx) error { 217 members, err := tx.ZRangeByScore(bucket, key, 5, 2, nil) 218 require.NoError(t, err) 219 require.Len(t, members, 4) 220 return nil 221 }) 222 require.NoError(t, err) 223 }) 224 } 225 func TestTx_ZPeekMin(t *testing.T) { 226 bucket := "bucket" 227 key := GetTestBytes(0) 228 229 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 230 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 231 232 for i := 0; i < 30; i++ { 233 txZAdd(t, db, bucket, key, GetTestBytes(i), float64(i), nil, nil) 234 } 235 236 // get minimum node 237 txZPeekMin(t, db, bucket, key, GetTestBytes(0), float64(0), nil, nil) 238 239 // bucket not exists 240 txZPeekMin(t, db, "non-exists-bucket", key, []byte{}, float64(0), ErrBucketNotExist, ErrBucketNotExist) 241 242 // key not exists 243 txZPeekMin(t, db, bucket, []byte("non-exists-key"), []byte{}, float64(0), ErrSortedSetNotFound, ErrSortedSetNotFound) 244 245 // add nodes 246 247 // add node that will not affect the minimum node 248 txZAdd(t, db, bucket, key, []byte("new-mem"), float64(3), nil, nil) 249 txZPeekMin(t, db, bucket, key, GetTestBytes(0), float64(0), nil, nil) 250 // add a new minimum value 251 txZAdd(t, db, bucket, key, []byte("new-min"), float64(0), nil, nil) 252 txZPeekMin(t, db, bucket, key, []byte("new-min"), float64(0), nil, nil) 253 254 // remove nodes 255 256 // remove minimum node 257 txZRem(t, db, bucket, key, []byte("new-min"), nil) 258 txZPeekMin(t, db, bucket, key, GetTestBytes(0), float64(0), nil, nil) 259 260 // remove non-minimum node 261 txZRem(t, db, bucket, key, GetTestBytes(5), nil) 262 txZPeekMin(t, db, bucket, key, GetTestBytes(0), float64(0), nil, nil) 263 264 // remove range by rank 265 err := db.Update(func(tx *Tx) error { 266 err := tx.ZRemRangeByRank(bucket, key, 1, 10) 267 assert.NoError(t, err) 268 return nil 269 }) 270 assert.NoError(t, err) 271 txZPeekMin(t, db, bucket, key, GetTestBytes(10), float64(10), nil, nil) 272 273 // pop 274 275 // pop min 276 txZPop(t, db, bucket, key, false, GetTestBytes(10), float64(10), nil) 277 txZPeekMin(t, db, bucket, key, GetTestBytes(11), float64(11), nil, nil) 278 279 // pop max 280 txZPop(t, db, bucket, key, true, GetTestBytes(29), float64(29), nil) 281 txZPeekMin(t, db, bucket, key, GetTestBytes(11), float64(11), nil, nil) 282 }) 283 } 284 285 func TestTx_ZRangeByRank(t *testing.T) { 286 bucket := "bucket" 287 key := GetTestBytes(0) 288 289 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 290 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 291 292 err := db.View(func(tx *Tx) error { 293 _, err := tx.ZRangeByRank(bucket, key, 1, 10) 294 require.Error(t, err) 295 return nil 296 }) 297 require.NoError(t, err) 298 299 for i := 0; i < 10; i++ { 300 txZAdd(t, db, bucket, key, GetTestBytes(i), float64(i), nil, nil) 301 } 302 303 err = db.View(func(tx *Tx) error { 304 members, err := tx.ZRangeByRank(bucket, key, 1, 10) 305 require.NoError(t, err) 306 require.Len(t, members, 10) 307 return nil 308 }) 309 require.NoError(t, err) 310 311 err = db.View(func(tx *Tx) error { 312 members, err := tx.ZRangeByRank(bucket, key, 3, 6) 313 require.NoError(t, err) 314 require.Len(t, members, 4) 315 return nil 316 }) 317 require.NoError(t, err) 318 319 err = db.View(func(tx *Tx) error { 320 members, err := tx.ZRangeByRank(bucket, key, -1, 11) 321 require.NoError(t, err) 322 require.Len(t, members, 1) 323 return nil 324 }) 325 require.NoError(t, err) 326 327 err = db.View(func(tx *Tx) error { 328 members, err := tx.ZRangeByRank(bucket, key, 8, 4) 329 require.NoError(t, err) 330 require.Len(t, members, 5) 331 332 for i, member := range members { 333 require.Equal(t, member.Score, float64(7-i)) 334 require.Equal(t, member.Value, GetTestBytes(7-i)) 335 } 336 337 return nil 338 }) 339 require.NoError(t, err) 340 }) 341 } 342 343 func TestTx_ZRemRangeByRank(t *testing.T) { 344 bucket := "bucket" 345 key := GetTestBytes(0) 346 347 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 348 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 349 350 err := db.Update(func(tx *Tx) error { 351 err := tx.ZRemRangeByRank(bucket, key, 1, 10) 352 assert.NoError(t, err) 353 return nil 354 }) 355 assert.NoError(t, err) 356 357 for i := 0; i < 10; i++ { 358 txZAdd(t, db, bucket, key, GetTestBytes(i), float64(i), nil, nil) 359 } 360 361 err = db.Update(func(tx *Tx) error { 362 err := tx.ZRemRangeByRank(bucket, key, 1, 10) 363 assert.NoError(t, err) 364 return nil 365 }) 366 assert.NoError(t, err) 367 368 txZCard(t, db, bucket, key, 0, nil) 369 370 for i := 0; i < 10; i++ { 371 txZAdd(t, db, bucket, key, GetTestBytes(i), float64(i), nil, nil) 372 } 373 374 err = db.Update(func(tx *Tx) error { 375 err := tx.ZRemRangeByRank(bucket, key, 1, 2) 376 assert.NoError(t, err) 377 return nil 378 }) 379 380 for i := 0; i < 2; i++ { 381 txZScore(t, db, bucket, key, GetTestBytes(0), 0, ErrSortedSetMemberNotExist) 382 } 383 384 err = db.Update(func(tx *Tx) error { 385 card, err := tx.ZCard(bucket, key) 386 assert.NoError(t, err) 387 assert.Equal(t, 8, card) 388 389 err = tx.ZRemRangeByRank(bucket, key, 6, 8) 390 assert.NoError(t, err) 391 return nil 392 }) 393 394 for i := 5; i < 8; i++ { 395 txZScore(t, db, bucket, key, GetTestBytes(0), 0, ErrSortedSetMemberNotExist) 396 } 397 398 err = db.Update(func(tx *Tx) error { 399 card, err := tx.ZCard(bucket, key) 400 assert.NoError(t, err) 401 assert.Equal(t, 5, card) 402 403 err = tx.ZRemRangeByRank(bucket, key, 4, 3) 404 assert.NoError(t, err) 405 return nil 406 }) 407 408 for i := 2; i < 4; i++ { 409 txZScore(t, db, bucket, key, GetTestBytes(0), 0, ErrSortedSetMemberNotExist) 410 } 411 412 assert.NoError(t, err) 413 }) 414 } 415 416 func TestTx_ZRank(t *testing.T) { 417 bucket := "bucket" 418 key := GetTestBytes(0) 419 420 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 421 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 422 423 txZRank(t, db, bucket, key, GetTestBytes(0), true, 0, ErrSortedSetNotFound) 424 txZRank(t, db, bucket, key, GetTestBytes(0), false, 0, ErrSortedSetNotFound) 425 426 for i := 0; i < 10; i++ { 427 txZAdd(t, db, bucket, key, GetTestBytes(i), float64(i), nil, nil) 428 } 429 430 txZRank(t, db, bucket, key, GetTestBytes(0), true, 10, nil) 431 txZRank(t, db, bucket, key, GetTestBytes(0), false, 1, nil) 432 433 txZRem(t, db, bucket, key, GetTestBytes(0), nil) 434 435 txZRank(t, db, bucket, key, GetTestBytes(0), true, 10, ErrSortedSetMemberNotExist) 436 txZRank(t, db, bucket, key, GetTestBytes(0), false, 1, ErrSortedSetMemberNotExist) 437 438 txZRem(t, db, bucket, key, GetTestBytes(3), nil) 439 440 txZRank(t, db, bucket, key, GetTestBytes(4), true, 6, nil) 441 txZRank(t, db, bucket, key, GetTestBytes(4), false, 3, nil) 442 }) 443 } 444 445 func TestTx_ZKeys(t *testing.T) { 446 bucket := "bucket" 447 key := "key_%d" 448 val := GetTestBytes(0) 449 450 runNutsDBTest(t, nil, func(t *testing.T, db *DB) { 451 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 452 for i := 0; i < 3; i++ { 453 txZAdd(t, db, bucket, []byte(fmt.Sprintf(key, i)), val, float64(i), nil, nil) 454 } 455 txZAdd(t, db, bucket, []byte("foo"), val, 1, nil, nil) 456 457 tests := []struct { 458 pattern string 459 expectedMatches int 460 expectedError error 461 }{ 462 {"*", 4, nil}, // find all keys 463 {"key_*", 3, nil}, // find keys with 'key_' prefix 464 {"fake_key*", 0, nil}, // find non-existing keys 465 } 466 467 for _, test := range tests { 468 txZKeys(t, db, bucket, test.pattern, func(key string) bool { return true }, test.expectedMatches, test.expectedError) 469 } 470 471 // stop after finding the expected number of keys. 472 expectNum := 2 473 var foundKeys []string 474 txZKeys(t, db, bucket, "*", func(key string) bool { 475 foundKeys = append(foundKeys, key) 476 return len(foundKeys) < expectNum 477 }, expectNum, nil) 478 assert.Equal(t, expectNum, len(foundKeys)) 479 }) 480 } 481 482 func TestTx_ZSetEntryIdxMode_HintKeyValAndRAMIdxMode(t *testing.T) { 483 bucket := "bucket" 484 key := GetTestBytes(0) 485 value := GetRandomBytes(24) 486 487 opts := DefaultOptions 488 opts.EntryIdxMode = HintKeyValAndRAMIdxMode 489 490 // HintKeyValAndRAMIdxMode 491 runNutsDBTest(t, &opts, func(t *testing.T, db *DB) { 492 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 493 494 err := db.Update(func(tx *Tx) error { 495 err := tx.ZAdd(bucket, key, float64(0), value) 496 require.NoError(t, err) 497 498 return nil 499 }) 500 require.NoError(t, err) 501 502 zset := db.Index.sortedSet.getWithDefault(1, db).M[string(key)] 503 hash, _ := getFnv32(value) 504 node := zset.dict[hash] 505 506 require.NotNil(t, node.record.Value) 507 require.Equal(t, value, node.record.Value) 508 }) 509 } 510 511 func TestTx_ZSetEntryIdxMode_HintKeyAndRAMIdxMode(t *testing.T) { 512 bucket := "bucket" 513 key := GetTestBytes(0) 514 value := GetRandomBytes(24) 515 516 opts := DefaultOptions 517 opts.EntryIdxMode = HintKeyAndRAMIdxMode 518 519 // HintKeyValAndRAMIdxMode 520 runNutsDBTest(t, &opts, func(t *testing.T, db *DB) { 521 txCreateBucket(t, db, DataStructureSortedSet, bucket, nil) 522 err := db.Update(func(tx *Tx) error { 523 err := tx.ZAdd(bucket, key, float64(0), value) 524 require.NoError(t, err) 525 526 return nil 527 }) 528 require.NoError(t, err) 529 530 zset := db.Index.sortedSet.getWithDefault(1, db).M[string(key)] 531 hash, _ := getFnv32(value) 532 node := zset.dict[hash] 533 534 require.Nil(t, node.record.Value) 535 536 v, err := db.getValueByRecord(node.record) 537 require.NoError(t, err) 538 require.Equal(t, value, v) 539 }) 540 }