diff --git a/cmd/index_analyzer/main.go b/cmd/index_analyzer/main.go index b7422da4..da158d0c 100644 --- a/cmd/index_analyzer/main.go +++ b/cmd/index_analyzer/main.go @@ -200,10 +200,19 @@ func analyzeIndex( logger.Fatal("error unpacking lids block", zap.Error(err)) } - last := len(block.Offsets) - 2 - for i := 0; i <= last; i++ { - tokenLIDs = append(tokenLIDs, block.LIDs[block.Offsets[i]:block.Offsets[i+1]]...) - if i < last || block.IsLastLID { // the end of token lids + listsCount := block.GetCount() + for i := 0; i < listsCount; i++ { + lidsBatch := block.GetLIDs(i) + iter := lidsBatch.Iter() + for { + lid, ok := iter.Next() + if !ok { + break + } + tokenLIDs = append(tokenLIDs, lid) + } + + if i < listsCount || block.IsLastLID() { // the end of token lids lidsTotal += len(tokenLIDs) lidsLens[tid] = len(tokenLIDs) lidsUniq[getLIDsHash(tokenLIDs)] = len(tokenLIDs) diff --git a/cmd/seq-db/seq-db.go b/cmd/seq-db/seq-db.go index 25b65f32..4467d0b5 100644 --- a/cmd/seq-db/seq-db.go +++ b/cmd/seq-db/seq-db.go @@ -273,6 +273,7 @@ func startStore( TokenTableZstdLevel: cfg.Compression.SealedZstdCompressionLevel, DocBlocksZstdLevel: cfg.Compression.DocBlockZstdCompressionLevel, DocBlockSize: int(cfg.DocsSorting.DocBlockSize), + LIDsBitmapThreshold: cfg.Sealing.Lids.BitmapThreshold, }, Fraction: frac.Config{ Search: frac.SearchConfig{ diff --git a/config/config.go b/config/config.go index 176d51a3..11faef42 100644 --- a/config/config.go +++ b/config/config.go @@ -75,6 +75,10 @@ type Config struct { Lids struct { // BlockSize sets max lids (postings) saved per LIDs block. BlockSize int `config:"block_size" default:"65536"` + // BitmapThreshold specifies minimum number of LIDs in the lid list + // which are serialized as bitmap. LIDs lists with more elements use bitmap encoding, + // while smaller lists use delta encoding. + BitmapThreshold int `config:"bitmap_threshold" default:"65536"` } `config:"lids"` } `config:"sealing"` diff --git a/config/frac_version.go b/config/frac_version.go index 73c3261a..5ef2b35d 100644 --- a/config/frac_version.go +++ b/config/frac_version.go @@ -21,6 +21,12 @@ const ( // BinaryDataV4 - delta bitpack encoded MIDs and LIDs BinaryDataV4 + + // BinaryDataV5 - token frequencies stored in token blocks for large tokens + BinaryDataV5 + + // BinaryDataV6 - bitmap for sufficiently large LID lists, mixed LIDs block format + BinaryDataV6 ) -const CurrentFracVersion = BinaryDataV4 +const CurrentFracVersion = BinaryDataV6 diff --git a/frac/common/seal_params.go b/frac/common/seal_params.go index 5a73de4e..99216d15 100644 --- a/frac/common/seal_params.go +++ b/frac/common/seal_params.go @@ -10,8 +10,9 @@ type SealParams struct { TokenListZstdLevel int DocsPositionsZstdLevel int TokenTableZstdLevel int - DocBlocksZstdLevel int // DocBlocksZstdLevel is the zstd compress level of each document block. - LIDBlockSize int - DocBlockSize int // DocBlockSize is decompressed payload size of document block. + DocBlocksZstdLevel int // DocBlocksZstdLevel is the zstd compress level of each document block. + LIDBlockSize int + LIDsBitmapThreshold int // LIDsBitmapThreshold is the minimum number of LIDs in the lid list to serialize as bitmap. + DocBlockSize int // DocBlockSize is decompressed payload size of document block. } diff --git a/frac/fraction_test.go b/frac/fraction_test.go index 180ec4e2..41916ece 100644 --- a/frac/fraction_test.go +++ b/frac/fraction_test.go @@ -104,6 +104,7 @@ func (s *FractionTestSuite) SetupTestCommon() { DocBlocksZstdLevel: 1, LIDBlockSize: 512, DocBlockSize: 128 * int(units.KiB), + LIDsBitmapThreshold: 25, } var err error @@ -1385,6 +1386,18 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { fromTime: fromTime, toTime: toTime, }, + { + name: "complex AND+OR", + query: "(service:gateway OR service:proxy OR service:scheduler) AND " + + "(message:request OR message:failed) AND (level:1 OR level:2 OR level:3)", + filter: func(doc *testDoc) bool { + return (doc.service == gateway || doc.service == proxy || doc.service == "scheduler") && + (strings.Contains(doc.message, "request") || strings.Contains(doc.message, "failed")) && + (doc.level >= 1 && doc.level <= 3) + }, + fromTime: fromTime, + toTime: toTime, + }, { name: "service:gateway AND NOT (message:request OR message:timed OR level:[0 to 3])", query: "service:gateway AND NOT (message:request OR message:timed OR level:[0 to 3])", diff --git a/frac/processor/search.go b/frac/processor/search.go index 61b33a9d..73f6f424 100644 --- a/frac/processor/search.go +++ b/frac/processor/search.go @@ -175,10 +175,10 @@ func IndexSearch( return qpr, nil } -func batcher(evalTree node.Node, buf []node.LID) func(need int) []node.LID { +func batcher(evalTree node.Node, buf []node.LID, desc bool) func(need int) []node.LID { if batchNode, ok := tryConvertToBatchedTree(evalTree); ok { return func(need int) []node.LID { - buf = batchNode.NextBatch().LIDs(buf[:0]) + buf = batchNode.NextBatch(need).CopyLIDs(desc, buf[:0]) if len(buf) > need { buf = buf[:need] } @@ -227,7 +227,7 @@ func iterateEvalTree( mids := buffers.mids rids := buffers.rids - batchedEvalTree := batcher(evalTree, buffers.lids) + batchedEvalTree := batcher(evalTree, buffers.lids, params.Order.IsDesc()) timerEval := sw.Timer("eval_tree_next") timerMID := sw.Timer("get_mid") @@ -353,9 +353,9 @@ func sampler(n uint32) func(in []node.LID) []node.LID { func tryConvertToBatchedTree(evalTree node.Node) (node.BatchedNode, bool) { switch it := evalTree.(type) { case *lids.IteratorDesc: - return it, true + return lids.NewBatchedIteratorDesc(it), true case *lids.IteratorAsc: - return it, true + return lids.NewBatchedIteratorAsc(it), true default: return nil, false } diff --git a/frac/sealed/lids/block.go b/frac/sealed/lids/block.go index 2f50e45d..b44221f5 100644 --- a/frac/sealed/lids/block.go +++ b/frac/sealed/lids/block.go @@ -1,53 +1,282 @@ package lids import ( + "bytes" + "encoding/binary" + "fmt" "math" + "slices" "unsafe" + "github.com/RoaringBitmap/roaring/v2" + "github.com/ozontech/seq-db/config" + "github.com/ozontech/seq-db/consts" + "github.com/ozontech/seq-db/node" "github.com/ozontech/seq-db/packer" ) -type Block struct { - LIDs []uint32 - Offsets []uint32 - // todo remove this legacy field +const ( + defaultLidsBitmapThreshold = math.MaxInt // bitmaps disabled by default +) + +type BlockPacker struct { + buf []uint32 // bitpack buffer (reusable across packing) + bmIndexes []uint32 // bitmap indexes (reusable across packing) + bitpackLIDs []uint32 + bitpackOffsets []uint32 + bm *roaring.Bitmap // reusable across packing + LidsBitmapThreshold int +} + +func NewBlockPacker() *BlockPacker { + return &BlockPacker{ + bm: roaring.NewBitmap(), + buf: make([]uint32, 0, consts.DefaultLIDBlockCap), + bitpackLIDs: make([]uint32, 0, consts.DefaultLIDBlockCap), + bitpackOffsets: make([]uint32, 0, consts.DefaultLIDBlockCap), + bmIndexes: make([]uint32, 0, consts.DefaultLIDBlockCap), + LidsBitmapThreshold: defaultLidsBitmapThreshold, + } +} + +// UnpackedBlock contains accumulated LIDs ready to pack. It's only used on sealing/compaction (index writing) time. +type UnpackedBlock struct { + LIDs []uint32 + Offsets []uint32 IsLastLID bool } -func (b *Block) getCount() int { - return len(b.Offsets) - 1 +// Block contains LIDs in variable format. It's used during search/queries processing. +// Field types is used to distinguish format for every LID list. If it's positive or zero, then it's a slot index in offsets +// (delta-encoding is used). If it's negative, then it's a bitmap slot index (index starts from -1, so -1 stands for bitmaps[0]). +// If types is nil, then the entire block is delta-encoded. +// +// On-disk format: +// +// [listsCount: uint32] — number of LID lists in the block +// [bitmapsCount: uint32] — number of lists stored as roaring bitmaps +// [bitmaps: bitmapsCount × roaring bitmap, serial format] +// [bitmapIndexes: delta-bitpack []uint32] — sorted list indices encoded as bitmaps +// [offsets: delta-bitpack []uint32] — slice boundaries in the delta-encoded LIDs array +// [lids: delta-bitpack []uint32] — concatenated delta-encoded LID values +// +// Each list i in [0, listsCount) is either a roaring bitmap (when i appears in bitmapIndexes) +// or a delta-encoded slice lids[offsets[k]:offsets[k+1]], where k is its delta-encoded slot index. +// Lists with length >= LidsBitmapThreshold are stored as bitmaps; shorter lists use delta-encoding. +type Block struct { + types []int32 // determines LID list type: delta-encoded (non-negative value) or bitmap (negative value). nil for delta-encoded blocks + lids []uint32 // all LIDs which are delta-encoded as a flat array + offsets []uint32 // offsets for delta-encoded LIDs + bitmaps []*roaring.Bitmap // all LIDs lists which are stored as bitmaps + lastLID bool // legacy field, will be removed soon +} + +func (b *Block) GetCount() int { + if b.types != nil { + return len(b.types) + } + + return len(b.offsets) - 1 } -func (b *Block) GetLIDs(i int) []uint32 { - return b.LIDs[b.Offsets[i]:b.Offsets[i+1]] +func (b *Block) IsLastLID() bool { + return b.lastLID } -func (b *Block) Pack(dst []byte, buf []uint32) []byte { - dst = packer.CompressDeltaBitpackUint32(dst, b.Offsets, buf) - dst = packer.CompressDeltaBitpackUint32(dst, b.LIDs, buf) - return dst +func (b *Block) GetLIDs(i int) node.LIDBatch { + if b.types == nil { + return node.NewSliceBatch(b.lids[b.offsets[i]:b.offsets[i+1]]) + } + t := b.types[i] + if t >= 0 { + return node.NewSliceBatch(b.lids[b.offsets[t]:b.offsets[t+1]]) + } + return node.NewBitmapBatch(b.bitmaps[-t-1]) +} + +func (b *Block) CopyLIDs(idx int, dst []uint32) []uint32 { + if b.types == nil { + dst = append(dst, b.lids[b.offsets[idx]:b.offsets[idx+1]]...) + return dst + } + t := b.types[idx] + if t >= 0 { + dst = append(dst, b.lids[b.offsets[t]:b.offsets[t+1]]...) + return dst + } + return b.copyLIDsFromBitmap(t, dst) +} + +func (b *Block) copyLIDsFromBitmap(ref int32, buf []uint32) []uint32 { + bitmap := b.bitmaps[-ref-1] + n := int(bitmap.GetCardinality()) + oldLen := len(buf) + + buf = slices.Grow(buf, n)[:oldLen+n] + dest := buf[oldLen:] + bitmap.ToExistingArray(&dest) + return buf } func (b *Block) GetSizeBytes() int { - const ( - uint32Size = int(unsafe.Sizeof(uint32(0))) - blockSize = int(unsafe.Sizeof(*b)) + const uint32Size = int(unsafe.Sizeof(uint32(0))) + size := int(unsafe.Sizeof(*b)) + uint32Size*cap(b.types) + uint32Size*cap(b.lids) + uint32Size*cap(b.offsets) + for _, bm := range b.bitmaps { + if bm != nil { + size += int(bm.GetSizeInBytes()) + } + } + return size +} + +func (p *BlockPacker) Pack(b *UnpackedBlock, dst []byte) []byte { + p.buf = p.buf[:0] + totalLists := len(b.Offsets) - 1 + bmCount := 0 // count of lid lists that will be stored as bitmaps + bmIndexes := p.bmIndexes[:0] + for i := 0; i < totalLists; i++ { + if int(b.Offsets[i+1]-b.Offsets[i]) >= p.LidsBitmapThreshold { + bmCount++ + bmIndexes = append(bmIndexes, uint32(i)) + } + } + + // write total number of LID lists and bitmap indexes + var numBuf [4]byte + binary.LittleEndian.PutUint32(numBuf[:], uint32(totalLists)) + dst = append(dst, numBuf[:]...) + binary.LittleEndian.PutUint32(numBuf[:], uint32(bmCount)) + dst = append(dst, numBuf[:]...) + + var ( + bitpackLIDs []uint32 + bitpackOffsets []uint32 ) - return blockSize + uint32Size*cap(b.LIDs) + uint32Size*cap(b.Offsets) + + if bmCount > 0 { + bitpackLIDs = p.bitpackLIDs[:0] + bitpackOffsets = p.bitpackOffsets[:0] + if bmCount < totalLists { + bitpackOffsets = append(bitpackOffsets, 0) + } + + for i := 0; i < totalLists; i++ { + lids := b.LIDs[b.Offsets[i]:b.Offsets[i+1]] + if len(lids) >= p.LidsBitmapThreshold { + dst = p.packBitmap(dst, lids) + } else { + bitpackLIDs = append(bitpackLIDs, lids...) + bitpackOffsets = append(bitpackOffsets, uint32(len(bitpackLIDs))) + } + } + } else { + bitpackLIDs = b.LIDs + bitpackOffsets = b.Offsets + } + + dst = packer.CompressDeltaBitpackUint32(dst, bmIndexes, p.buf) + dst = packer.CompressDeltaBitpackUint32(dst, bitpackOffsets, p.buf) + dst = packer.CompressDeltaBitpackUint32(dst, bitpackLIDs, p.buf) + return dst +} + +func (p *BlockPacker) packBitmap(dst []byte, lids []uint32) []byte { + p.bm.Clear() + p.bm.AddMany(lids) + p.bm.RunOptimize() + + wrt := bytes.NewBuffer(dst) + _, err := p.bm.WriteTo(wrt) + if err != nil { + panic(fmt.Errorf("bitmap write failed: %w", err)) + } + return wrt.Bytes() } func (b *Block) Unpack(data []byte, fracVer config.BinaryDataVersion, buf *UnpackBuffer) error { buf.Reset(fracVer) - if fracVer >= config.BinaryDataV4 { - return b.unpackBitpack(data, buf) + if fracVer < config.BinaryDataV4 { + return b.unpackVarintsV1(data, buf) + } + if fracVer < config.BinaryDataV6 { + return b.unpackBitpackV4(data, buf) + } + + return b.unpackBlockV6(data, buf) +} + +// unpackBlockV6 unpacks the mixed bitmap / delta-bitpack format (BinaryDataV6+). +func (b *Block) unpackBlockV6(data []byte, buf *UnpackBuffer) error { + listsCount := int(binary.LittleEndian.Uint32(data[:4])) + data = data[4:] + bitmapsCount := int(binary.LittleEndian.Uint32(data[:4])) + data = data[4:] + + bitmaps := make([]*roaring.Bitmap, bitmapsCount) + for i := 0; i < bitmapsCount; i++ { + rb := roaring.NewBitmap() + n, err := rb.ReadFrom(bytes.NewReader(data)) + if err != nil { + return fmt.Errorf("read bitmap %d: %w", i, err) + } + data = data[n:] + bitmaps[i] = rb + } + + var ( + err error + bitmapIndexes []uint32 + ) + data, bitmapIndexes, err = packer.DecompressDeltaBitpackUint32(data, buf.decompressed, buf.compressed) + if err != nil { + return err + } + b.types = deriveTypes(listsCount, bitmapIndexes) + + var values []uint32 + data, values, err = packer.DecompressDeltaBitpackUint32(data, buf.decompressed, buf.compressed) + if err != nil { + return err + } + offsets := append([]uint32{}, values...) + + _, values, err = packer.DecompressDeltaBitpackUint32(data, buf.decompressed, buf.compressed) + if err != nil { + return err } + lids := append([]uint32{}, values...) - return b.unpackVarint(data, buf) + b.lids = lids + b.offsets = offsets + b.bitmaps = bitmaps + return nil } -func (b *Block) unpackBitpack(data []byte, buf *UnpackBuffer) error { +// deriveTypes derives types array (only for LID blocks which have at least one bitmap). +func deriveTypes(totalLists int, bitmapIndexes []uint32) []int32 { + if len(bitmapIndexes) == 0 { + return nil + } + listTypes := make([]int32, totalLists) + bmIdx := 0 + deltaIdx := 0 + bmSlotIdx := 0 + for i := 0; i < totalLists; i++ { + if bmSlotIdx < len(bitmapIndexes) && bitmapIndexes[bmSlotIdx] == uint32(i) { + listTypes[i] = -int32(bmIdx + 1) + bmIdx++ + bmSlotIdx++ + } else { + listTypes[i] = int32(deltaIdx) + deltaIdx++ + } + } + return listTypes +} + +func (b *Block) unpackBitpackV4(data []byte, buf *UnpackBuffer) error { var err error var values []uint32 @@ -55,22 +284,26 @@ func (b *Block) unpackBitpack(data []byte, buf *UnpackBuffer) error { if err != nil { return err } - b.Offsets = append([]uint32{}, values...) + offsets := append([]uint32{}, values...) _, values, err = packer.DecompressDeltaBitpackUint32(data, buf.decompressed, buf.compressed) if err != nil { return err } - b.LIDs = append([]uint32{}, values...) + lids := append([]uint32{}, values...) + + b.types = nil + b.lids = lids + b.offsets = offsets + b.bitmaps = nil + b.lastLID = false return nil } -func (b *Block) unpackVarint(data []byte, buf *UnpackBuffer) error { +func (b *Block) unpackVarintsV1(data []byte, buf *UnpackBuffer) error { var lid, offset uint32 - b.IsLastLID = true - - buf.offsets = append(buf.offsets, 0) // first offset is always zero + buf.offsets = append(buf.offsets, 0) unpacker := packer.NewBytesUnpacker(data) for unpacker.Len() > 0 { @@ -80,7 +313,7 @@ func (b *Block) unpackVarint(data []byte, buf *UnpackBuffer) error { } lid += uint32(delta) - if lid == math.MaxUint32 { // end of LIDs of current TID, see `Block.Pack()` method + if lid == math.MaxUint32 { offset = uint32(len(buf.lids)) buf.offsets = append(buf.offsets, offset) lid -= uint32(delta) @@ -90,14 +323,16 @@ func (b *Block) unpackVarint(data []byte, buf *UnpackBuffer) error { buf.lids = append(buf.lids, lid) } + lastLID := true if int(offset) < len(buf.lids) { - b.IsLastLID = false + lastLID = false buf.offsets = append(buf.offsets, uint32(len(buf.lids))) } - // copy from buffer - b.LIDs = append([]uint32{}, buf.lids...) - b.Offsets = append([]uint32{}, buf.offsets...) - + b.types = nil + b.lids = append([]uint32{}, buf.lids...) + b.offsets = append([]uint32{}, buf.offsets...) + b.bitmaps = nil + b.lastLID = lastLID return nil } diff --git a/frac/sealed/lids/block_test.go b/frac/sealed/lids/block_test.go index fde6d26f..d2f3d9d6 100644 --- a/frac/sealed/lids/block_test.go +++ b/frac/sealed/lids/block_test.go @@ -9,6 +9,7 @@ import ( "github.com/stretchr/testify/require" "github.com/ozontech/seq-db/config" + "github.com/ozontech/seq-db/node" ) func TestBlockPack(t *testing.T) { @@ -44,12 +45,12 @@ func TestBlockPack(t *testing.T) { offsets: []uint32{0, 3, 6, 8}, }, { - name: "medium_many_tokens", + name: "large_delta_encoded", generator: func() ([]uint32, []uint32) { lids := make([]uint32, 0) offsets := []uint32{0} startLID := uint32(100) - for i := 0; i < 10; i++ { + for i := 0; i < 100; i++ { for j := 0; j < 3; j++ { lids = append(lids, startLID+uint32(i*10+j)) } @@ -90,17 +91,32 @@ func TestBlockPack(t *testing.T) { offsets: []uint32{0, 129}, }, { - name: "medium_4k_lids", + name: "medium_4k_bitmap_and_small_list", + lids: generate(4096), + offsets: []uint32{0, 4093, 4096}, + }, + { + name: "medium_4k_small_list_and_bitmap", + lids: generate(4096), + offsets: []uint32{0, 3, 4096}, + }, + { + name: "medium_4k_hybrid", + lids: generate(4096), + offsets: []uint32{0, 1000, 1005, 1010, 2000, 2100, 2103, 2106, 2107, 3000, 3500, 3505, 4096}, + }, + { + name: "medium_4k", lids: generate(4096), offsets: []uint32{0, 4096}, }, { - name: "medium_4k_minus_one_lids", + name: "medium_4095", lids: generate(4095), offsets: []uint32{0, 10, 50, 100, 150, 190, 1000, 1500, 4095}, }, { - name: "medium_4k_plus_one_lids", + name: "medium_4097", lids: generate(4097), offsets: []uint32{0, 10, 50, 100, 150, 190, 1000, 1500, 4097}, }, @@ -123,104 +139,137 @@ func TestBlockPack(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { - var lids []uint32 + var lidsIn []uint32 var offsets []uint32 if tc.generator != nil { - lids, offsets = tc.generator() + lidsIn, offsets = tc.generator() } else { - lids = tc.lids + lidsIn = tc.lids offsets = tc.offsets } - block := &Block{ - LIDs: lids, + block := &UnpackedBlock{ + LIDs: lidsIn, Offsets: offsets, } - packed := block.Pack(nil, nil) + packer := NewBlockPacker() + packer.LidsBitmapThreshold = 25 + + packed := packer.Pack(block, nil) require.NotEmpty(t, packed) - unpacked := &Block{} buf := &UnpackBuffer{} - err := unpacked.Unpack(packed, config.CurrentFracVersion, buf) + var unpacked Block + require.NoError(t, unpacked.Unpack(packed, config.CurrentFracVersion, buf)) - require.NoError(t, err) - assert.EqualExportedValues(t, block, unpacked) + assertListsEqual(t, block, &unpacked) }) } } -func generate(n int) []uint32 { - v := make([]uint32, n) - last := uint32(100) - for i := range v { - v[i] = last - last += uint32(1 + rand.Intn(5)) +func TestBlockPack_VariableMixed(t *testing.T) { + small := generate(10) + large := generate(30) + block := &UnpackedBlock{ + LIDs: append(append([]uint32{}, small...), large...), + Offsets: []uint32{0, uint32(len(small)), uint32(len(small) + len(large))}, } - return v + + packed := NewBlockPacker().Pack(block, nil) + + buf := &UnpackBuffer{} + var ub Block + require.NoError(t, ub.Unpack(packed, config.CurrentFracVersion, buf)) + assert.Equal(t, 2, ub.GetCount()) + assert.Equal(t, small, ToArray(ub.GetLIDs(0))) + assert.Equal(t, large, ToArray(ub.GetLIDs(1))) +} + +func ToArray(b node.LIDBatch) []uint32 { + if b.IsEmpty() { + return nil + } + out := make([]uint32, 0, b.Len()) + for _, lid := range b.CopyLIDs(true, nil) { + out = append(out, lid.Unpack()) + } + return out } func TestBlockPack_ReuseBuffer(t *testing.T) { - // Test that UnpackBuffer can be reused - block1 := &Block{ + block1 := &UnpackedBlock{ LIDs: generate(64 * 1024), Offsets: []uint32{0, 3}, } - block2 := &Block{ + block2 := &UnpackedBlock{ LIDs: generate(64 * 1024), Offsets: []uint32{0, 4}, } - buf1 := make([]uint32, 0, 64*1024) - packed1 := block1.Pack(nil, buf1) - - buf1 = buf1[:0] - packed2 := block2.Pack(nil, buf1) + packer := NewBlockPacker() + packed1 := packer.Pack(block1, nil) + packed2 := packer.Pack(block2, nil) buf2 := &UnpackBuffer{} - unpacked1 := &Block{} - err := unpacked1.Unpack(packed1, config.CurrentFracVersion, buf2) - require.NoError(t, err) - assert.Equal(t, block1.LIDs, unpacked1.LIDs) + var unpacked1, unpacked2 Block + require.NoError(t, unpacked1.Unpack(packed1, config.CurrentFracVersion, buf2)) + require.NoError(t, unpacked2.Unpack(packed2, config.CurrentFracVersion, buf2)) + + assertListsEqual(t, block1, &unpacked1) + assertListsEqual(t, block2, &unpacked2) +} - unpacked2 := &Block{} - err = unpacked2.Unpack(packed2, config.CurrentFracVersion, buf2) - require.NoError(t, err) - assert.Equal(t, block2.LIDs, unpacked2.LIDs) +func assertListsEqual(t *testing.T, src *UnpackedBlock, blk *Block) { + t.Helper() + require.Equal(t, len(src.Offsets)-1, blk.GetCount()) + for i := 0; i < blk.GetCount(); i++ { + want := src.LIDs[src.Offsets[i]:src.Offsets[i+1]] + assert.Equal(t, want, ToArray(blk.GetLIDs(i))) + } +} + +func generate(n int) []uint32 { + v := make([]uint32, n) + last := uint32(100) + for i := range v { + v[i] = last + last += uint32(1 + rand.Intn(5)) + } + return v } func BenchmarkBlock_Pack(b *testing.B) { - lids := generate(64 * 1024) + lidsIn := generate(64 * 1024) - block := &Block{ - LIDs: lids, + block := &UnpackedBlock{ + LIDs: lidsIn, Offsets: []uint32{0, 64 * 1024}, } - tmp := make([]uint32, 0, 64*1024/4) + packer := NewBlockPacker() for b.Loop() { - block.Pack(nil, tmp) + packer.Pack(block, nil) } } func BenchmarkBlock_Unpack(b *testing.B) { - lids := generate(64 * 1024) + lidsIn := generate(64 * 1024) - block := &Block{ - LIDs: lids, + block := &UnpackedBlock{ + LIDs: lidsIn, Offsets: []uint32{0, 64 * 1024}, } - packed := block.Pack(nil, nil) + packed := NewBlockPacker().Pack(block, nil) buf := &UnpackBuffer{} - unpacked := &Block{} b.ResetTimer() for b.Loop() { - err := unpacked.Unpack(packed, config.CurrentFracVersion, buf) - assert.NoError(b, err) + var ub Block + assert.NoError(b, ub.Unpack(packed, config.CurrentFracVersion, buf)) } } diff --git a/frac/sealed/lids/cursor.go b/frac/sealed/lids/cursor.go index 4dc0d73c..d3446ca7 100644 --- a/frac/sealed/lids/cursor.go +++ b/frac/sealed/lids/cursor.go @@ -1,5 +1,7 @@ package lids +import "github.com/ozontech/seq-db/node" + type Counter interface { AddLIDsCount(int) } @@ -15,7 +17,7 @@ type Cursor struct { blockIndex uint32 tryNextBlock bool - lids []uint32 + batch node.LIDBatch counter Counter } @@ -38,6 +40,7 @@ func NewLIDsCursor( tid: tid, blockIndex: startIndex, tryNextBlock: true, + batch: node.EmptyBatch(), counter: counter, } diff --git a/frac/sealed/lids/iterator_asc.go b/frac/sealed/lids/iterator_asc.go index 5f7a5f03..6f73a167 100644 --- a/frac/sealed/lids/iterator_asc.go +++ b/frac/sealed/lids/iterator_asc.go @@ -1,44 +1,61 @@ package lids import ( - "sort" - "go.uber.org/zap" "github.com/ozontech/seq-db/logger" "github.com/ozontech/seq-db/node" ) -type IteratorAsc Cursor +type IteratorAsc struct { + Cursor + it node.Iter +} + +func NewIteratorAsc( + table *Table, + loader *Loader, + startIndex uint32, + tid uint32, + counter Counter, + minLID, maxLID uint32, +) *IteratorAsc { + it := &IteratorAsc{ + Cursor: *NewLIDsCursor(table, loader, startIndex, tid, counter, minLID, maxLID), + } + it.it = it.batch.ReverseIter() + return it +} func (*IteratorAsc) String() string { return "LIDS_ASC" } -// narrowLIDsRange cuts LIDs between from and to. Returns new lids and tryNextBlock flag -func (it *IteratorAsc) narrowLIDsRange(lids []uint32, tryNextBlock bool) ([]uint32, bool) { - first := lids[0] - if it.maxLID < first { // fast path: out-of-bounds 1; allowed to continue reading blocks - return nil, tryNextBlock +// narrowLIDsRange cuts LIDs between minLID and maxLID. Returns updated tryNextBlock flag. +func (it *IteratorAsc) narrowLIDsRange(tryNextBlock bool) bool { + if it.batch.IsEmpty() { + return tryNextBlock } - last := lids[len(lids)-1] - if it.minLID > last { // fast path: out-of-bounds 2 - return nil, false // stop reading blocks + first := it.batch.Min() + if it.maxLID < first { + it.batch = node.EmptyBatch() + return tryNextBlock } - if it.minLID > first { - left := sort.Search(len(lids), func(i int) bool { return lids[i] >= it.minLID }) - lids = lids[left:] - tryNextBlock = false + last := it.batch.Max() + if it.minLID > last { + it.batch = node.EmptyBatch() + return false } - if it.maxLID <= last { - right := sort.Search(len(lids), func(i int) bool { return lids[i] > it.maxLID }) - lids = lids[:right] + lastBlock := it.minLID > first + it.batch = it.batch.Narrow(it.minLID, it.maxLID) + if lastBlock { + tryNextBlock = false } - return lids, tryNextBlock + return tryNextBlock } func (it *IteratorAsc) loadNextLIDsBlock() { @@ -47,92 +64,41 @@ func (it *IteratorAsc) loadNextLIDsBlock() { logger.Panic("error loading LIDs block", zap.Error(err)) } - if block.getCount() != int(it.table.GetChunksCount(it.blockIndex)) { + if block.GetCount() != int(it.table.GetChunksCount(it.blockIndex)) { logger.Panic("unexpected LIDs count") } - it.lids = block.GetLIDs(it.table.GetChunkIndex(it.blockIndex, it.tid)) - it.tryNextBlock = it.table.HasTIDInPrevBlock(it.blockIndex, it.tid) + it.batch = block.GetLIDs(it.table.GetChunkIndex(it.blockIndex, it.tid)) + tryNextBlock := it.table.HasTIDInPrevBlock(it.blockIndex, it.tid) + it.tryNextBlock = it.narrowLIDsRange(tryNextBlock) + it.it = it.batch.ReverseIter() + it.counter.AddLIDsCount(it.batch.Len()) it.blockIndex-- } func (it *IteratorAsc) Next() node.LID { - for len(it.lids) == 0 { + for { + lid, ok := it.it.Next() + if ok { + return node.NewAscLID(lid) + } if !it.tryNextBlock { return node.NullLID() } - - it.loadNextLIDsBlock() // last chunk in block but not last for tid; need load next block - it.lids, it.tryNextBlock = it.narrowLIDsRange(it.lids, it.tryNextBlock) - it.counter.AddLIDsCount(len(it.lids)) // inc loaded LIDs count + it.loadNextLIDsBlock() } - - i := len(it.lids) - 1 - lid := it.lids[i] - it.lids = it.lids[:i] - return node.NewAscLID(lid) } // NextGeq returns the next (in reverse iteration order) LID that is <= maxLID. func (it *IteratorAsc) NextGeq(nextID node.LID) node.LID { for { - for len(it.lids) == 0 { - if !it.tryNextBlock { - return node.NullLID() - } - - it.loadNextLIDsBlock() - it.lids, it.tryNextBlock = it.narrowLIDsRange(it.lids, it.tryNextBlock) - it.counter.AddLIDsCount(len(it.lids)) - } - - // fast path: smallest remaining > nextID => skip entire block - // TODO(cheb0): We could also pass LID into narrowLIDsRange to perform block skipping once we add something like MinLID to LID block header - if it.lids[0] > nextID.Unpack() { - it.lids = it.lids[:0] - continue - } - - idx := sort.Search(len(it.lids), func(i int) bool { return it.lids[i] > nextID.Unpack() }) - 1 - if idx >= 0 { - lid := it.lids[idx] - it.lids = it.lids[:idx] + lid, ok := it.it.NextGeq(nextID.Unpack()) + if ok { return node.NewAscLID(lid) } - - it.lids = it.lids[:0] - } -} - -func (it *IteratorAsc) NextBatch() node.LIDBatch { - return it.NextBatchGeq(node.NewAscZeroLID()) -} - -func (it *IteratorAsc) NextBatchGeq(nextID node.LID) node.LIDBatch { - for { - for len(it.lids) == 0 { - if !it.tryNextBlock { - return node.NewAscBatch(nil) - } - it.loadNextLIDsBlock() - it.lids, it.tryNextBlock = it.narrowLIDsRange(it.lids, it.tryNextBlock) - it.counter.AddLIDsCount(len(it.lids)) - } - - // fast path: smallest remaining > nextID => skip entire block - // TODO(cheb0): We could also pass LID into narrowLIDsRange to perform block skipping once we add something like MinLID to LID block header - if it.lids[0] > nextID.Unpack() { - it.lids = it.lids[:0] - continue - } - - idx := sort.Search(len(it.lids), func(i int) bool { return it.lids[i] > nextID.Unpack() }) - 1 - if idx >= 0 { - batch := it.lids[:idx+1] - it.lids = it.lids[:0] - return node.NewAscBatch(batch) + if !it.tryNextBlock { + return node.NullLID() } - - it.lids = it.lids[:0] + it.loadNextLIDsBlock() } } diff --git a/frac/sealed/lids/iterator_batched_asc.go b/frac/sealed/lids/iterator_batched_asc.go new file mode 100644 index 00000000..da9efed5 --- /dev/null +++ b/frac/sealed/lids/iterator_batched_asc.go @@ -0,0 +1,94 @@ +package lids + +import ( + "go.uber.org/zap" + + "github.com/ozontech/seq-db/logger" + "github.com/ozontech/seq-db/node" +) + +type BatchedIteratorAsc struct { + Cursor +} + +func NewBatchedIteratorAsc(it *IteratorAsc) *BatchedIteratorAsc { + return &BatchedIteratorAsc{ + Cursor: it.Cursor, + } +} + +func (*BatchedIteratorAsc) String() string { + return "LIDS_ASC_BATCHED" +} + +// narrowLIDsRange cuts LIDs between minLID and maxLID. Returns updated tryNextBlock flag. +func (it *BatchedIteratorAsc) narrowLIDsRange(tryNextBlock bool) bool { + if it.batch.IsEmpty() { + return tryNextBlock + } + + first := it.batch.Min() + if it.maxLID < first { + it.batch = node.EmptyBatch() + return tryNextBlock + } + + batchMax := it.batch.Max() + if it.minLID > batchMax { + it.batch = node.EmptyBatch() + return false + } + + lastBlock := it.minLID > first + it.batch = it.batch.Narrow(it.minLID, it.maxLID) + if lastBlock { + tryNextBlock = false + } + + return tryNextBlock +} + +func (it *BatchedIteratorAsc) loadNextLIDsBlock() { + block, err := it.loader.GetLIDsBlock(it.table.StartBlockIndex + it.blockIndex) + if err != nil { + logger.Panic("error loading LIDs block", zap.Error(err)) + } + + if block.GetCount() != int(it.table.GetChunksCount(it.blockIndex)) { + logger.Panic("unexpected LIDs count") + } + + it.batch = block.GetLIDs(it.table.GetChunkIndex(it.blockIndex, it.tid)) + it.tryNextBlock = it.table.HasTIDInPrevBlock(it.blockIndex, it.tid) + it.blockIndex-- +} + +func (it *BatchedIteratorAsc) NextBatch(need int) node.LIDBatch { + return it.NextBatchGeq(need, node.NewAscZeroLID()) +} + +func (it *BatchedIteratorAsc) NextBatchGeq(_ int, nextID node.LID) node.LIDBatch { + for { + if it.batch.IsEmpty() { + if !it.tryNextBlock { + return node.EmptyBatch() + } + it.loadNextLIDsBlock() + it.tryNextBlock = it.narrowLIDsRange(it.tryNextBlock) + it.counter.AddLIDsCount(it.batch.Len()) + } + + if it.batch.IsEmpty() { + continue + } + + if nextID.Unpack() < it.batch.Min() { + it.batch = node.EmptyBatch() + continue + } + + out := it.batch + it.batch = node.EmptyBatch() + return out + } +} diff --git a/frac/sealed/lids/iterator_batched_desc.go b/frac/sealed/lids/iterator_batched_desc.go new file mode 100644 index 00000000..318e312c --- /dev/null +++ b/frac/sealed/lids/iterator_batched_desc.go @@ -0,0 +1,94 @@ +package lids + +import ( + "go.uber.org/zap" + + "github.com/ozontech/seq-db/logger" + "github.com/ozontech/seq-db/node" +) + +type BatchedIteratorDesc struct { + Cursor +} + +func NewBatchedIteratorDesc(it *IteratorDesc) *BatchedIteratorDesc { + return &BatchedIteratorDesc{ + Cursor: it.Cursor, + } +} + +func (*BatchedIteratorDesc) String() string { + return "LIDS_DESC_BATCHED" +} + +// narrowLIDsRange cuts LIDs between minLID and maxLID. Returns updated tryNextBlock flag. +func (it *BatchedIteratorDesc) narrowLIDsRange(tryNextBlock bool) bool { + if it.batch.IsEmpty() { + return tryNextBlock + } + + first := it.batch.Min() + if it.maxLID < first { + it.batch = node.EmptyBatch() + return false + } + + batchMax := it.batch.Max() + if it.minLID > batchMax { + it.batch = node.EmptyBatch() + return tryNextBlock + } + + lastBlock := it.maxLID < batchMax + it.batch = it.batch.Narrow(it.minLID, it.maxLID) + if lastBlock { + tryNextBlock = false + } + + return tryNextBlock +} + +func (it *BatchedIteratorDesc) loadNextLIDsBlock() { + block, err := it.loader.GetLIDsBlock(it.table.StartBlockIndex + it.blockIndex) + if err != nil { + logger.Panic("error loading LIDs block", zap.Error(err)) + } + + if block.GetCount() != int(it.table.GetChunksCount(it.blockIndex)) { + logger.Panic("unexpected LIDs count") + } + + it.batch = block.GetLIDs(it.table.GetChunkIndex(it.blockIndex, it.tid)) + it.tryNextBlock = it.table.HasTIDInNextBlock(it.blockIndex, it.tid) + it.blockIndex++ +} + +func (it *BatchedIteratorDesc) NextBatch(need int) node.LIDBatch { + return it.NextBatchGeq(need, node.NewDescZeroLID()) +} + +func (it *BatchedIteratorDesc) NextBatchGeq(_ int, nextID node.LID) node.LIDBatch { + for { + if it.batch.IsEmpty() { + if !it.tryNextBlock { + return node.EmptyBatch() + } + it.loadNextLIDsBlock() + it.tryNextBlock = it.narrowLIDsRange(it.tryNextBlock) + it.counter.AddLIDsCount(it.batch.Len()) + } + + if it.batch.IsEmpty() { + continue + } + + if nextID.Unpack() > it.batch.Max() { + it.batch = node.EmptyBatch() + continue + } + + out := it.batch + it.batch = node.EmptyBatch() + return out + } +} diff --git a/frac/sealed/lids/iterator_desc.go b/frac/sealed/lids/iterator_desc.go index cb2551f4..b966e372 100644 --- a/frac/sealed/lids/iterator_desc.go +++ b/frac/sealed/lids/iterator_desc.go @@ -1,44 +1,61 @@ package lids import ( - "sort" - "go.uber.org/zap" "github.com/ozontech/seq-db/logger" "github.com/ozontech/seq-db/node" ) -type IteratorDesc Cursor +type IteratorDesc struct { + Cursor + it node.Iter +} + +func NewIteratorDesc( + table *Table, + loader *Loader, + startIndex uint32, + tid uint32, + counter Counter, + minLID, maxLID uint32, +) *IteratorDesc { + it := &IteratorDesc{ + Cursor: *NewLIDsCursor(table, loader, startIndex, tid, counter, minLID, maxLID), + } + it.it = it.batch.Iter() + return it +} func (*IteratorDesc) String() string { return "LIDS_DESC" } -// narrowLIDsRange cuts LIDs between from and to. Returns new lids and tryNextBlock flag -func (it *IteratorDesc) narrowLIDsRange(lids []uint32, tryNextBlock bool) ([]uint32, bool) { - first := lids[0] - if it.maxLID < first { // fast path: out-of-bounds 1 - return nil, false // stop reading blocks +// narrowLIDsRange cuts LIDs between minLID and maxLID. Returns updated tryNextBlock flag. +func (it *IteratorDesc) narrowLIDsRange(tryNextBlock bool) bool { + if it.batch.IsEmpty() { + return tryNextBlock } - last := lids[len(lids)-1] - if it.minLID > last { // fast path: out-of-bounds 2; allowed to continue reading blocks - return nil, tryNextBlock + first := it.batch.Min() + if it.maxLID < first { + it.batch = node.EmptyBatch() + return false } - if it.minLID > first { - left := sort.Search(len(lids), func(i int) bool { return lids[i] >= it.minLID }) - lids = lids[left:] + last := it.batch.Max() + if it.minLID > last { + it.batch = node.EmptyBatch() + return tryNextBlock } - if it.maxLID <= last { - right := sort.Search(len(lids), func(i int) bool { return lids[i] > it.maxLID }) - lids = lids[:right] + lastBlock := it.maxLID < last + it.batch = it.batch.Narrow(it.minLID, it.maxLID) + if lastBlock { tryNextBlock = false } - return lids, tryNextBlock + return tryNextBlock } func (it *IteratorDesc) loadNextLIDsBlock() { @@ -47,96 +64,41 @@ func (it *IteratorDesc) loadNextLIDsBlock() { logger.Panic("error loading LIDs block", zap.Error(err)) } - if block.getCount() != int(it.table.GetChunksCount(it.blockIndex)) { + if block.GetCount() != int(it.table.GetChunksCount(it.blockIndex)) { logger.Panic("unexpected LIDs count") } - it.lids = block.GetLIDs(it.table.GetChunkIndex(it.blockIndex, it.tid)) - it.tryNextBlock = it.table.HasTIDInNextBlock(it.blockIndex, it.tid) + it.batch = block.GetLIDs(it.table.GetChunkIndex(it.blockIndex, it.tid)) + it.counter.AddLIDsCount(it.batch.Len()) + tryNextBlock := it.table.HasTIDInNextBlock(it.blockIndex, it.tid) + it.tryNextBlock = it.narrowLIDsRange(tryNextBlock) + it.it = it.batch.Iter() it.blockIndex++ } func (it *IteratorDesc) Next() node.LID { - for len(it.lids) == 0 { + for { + v, ok := it.it.Next() + if ok { + return node.NewDescLID(v) + } if !it.tryNextBlock { return node.NullLID() } - - it.loadNextLIDsBlock() // last chunk in block but not last for tid; need load next block - it.lids, it.tryNextBlock = it.narrowLIDsRange(it.lids, it.tryNextBlock) - it.counter.AddLIDsCount(len(it.lids)) // inc loaded LIDs count + it.loadNextLIDsBlock() } - - lid := it.lids[0] - it.lids = it.lids[1:] - return node.NewDescLID(lid) } // NextGeq finds next greater or equal func (it *IteratorDesc) NextGeq(nextID node.LID) node.LID { for { - for len(it.lids) == 0 { - if !it.tryNextBlock { - return node.NullLID() - } - - it.loadNextLIDsBlock() // last chunk in block but not last for tid; need load next block - it.lids, it.tryNextBlock = it.narrowLIDsRange(it.lids, it.tryNextBlock) - it.counter.AddLIDsCount(len(it.lids)) // inc loaded LIDs count - } - - // fast path: last LID < nextID => skip the entire block - // TODO(cheb0): We could also pass LID into narrowLIDsRange to perform block skipping once we add something like MinLID to LID block header - if nextID.Unpack() > it.lids[len(it.lids)-1] { - it.lids = it.lids[:0] - continue + v, ok := it.it.NextGeq(nextID.Unpack()) + if ok { + return node.NewDescLID(v) } - - idx := sort.Search(len(it.lids), func(i int) bool { return it.lids[i] >= nextID.Unpack() }) - if idx < len(it.lids) { - it.lids = it.lids[idx:] - lid := it.lids[0] - it.lids = it.lids[1:] - return node.NewDescLID(lid) - } - - it.lids = it.lids[:0] - } -} - -func (it *IteratorDesc) NextBatch() node.LIDBatch { - return it.NextBatchGeq(node.NewDescZeroLID()) -} - -func (it *IteratorDesc) NextBatchGeq(nextID node.LID) node.LIDBatch { - for { - for len(it.lids) == 0 { - if !it.tryNextBlock { - return node.NewDescBatch(nil) - } - it.loadNextLIDsBlock() - it.lids, it.tryNextBlock = it.narrowLIDsRange(it.lids, it.tryNextBlock) - it.counter.AddLIDsCount(len(it.lids)) - } - last := it.lids[len(it.lids)-1] - if nextID.Unpack() > last { - it.lids = it.lids[:0] - continue - } - - // fast path: last LID < nextLID => skip the entire block - if nextID.Unpack() > it.lids[len(it.lids)-1] { - it.lids = it.lids[:0] - continue - } - - idx := sort.Search(len(it.lids), func(i int) bool { return it.lids[i] >= nextID.Unpack() }) - if idx < len(it.lids) { - batch := it.lids[idx:len(it.lids)] - it.lids = it.lids[:0] - return node.NewDescBatch(batch) + if !it.tryNextBlock { + return node.NullLID() } - - it.lids = it.lids[:0] + it.loadNextLIDsBlock() } } diff --git a/frac/sealed/lids/loader.go b/frac/sealed/lids/loader.go index cf987a97..37e50c45 100644 --- a/frac/sealed/lids/loader.go +++ b/frac/sealed/lids/loader.go @@ -80,10 +80,8 @@ func (l *Loader) readLIDsBlock(blockIndex uint32) (*Block, error) { } block := &Block{} - err = block.Unpack(l.blockBuf, l.fracVer, l.unpackBuf) - if err != nil { + if err := block.Unpack(l.blockBuf, l.fracVer, l.unpackBuf); err != nil { return nil, err } - - return block, err + return block, nil } diff --git a/frac/sealed_index.go b/frac/sealed_index.go index 84d563c3..cbc4fae3 100644 --- a/frac/sealed_index.go +++ b/frac/sealed_index.go @@ -264,12 +264,12 @@ func (ti *sealedTokenIndex) GetLIDsFromTIDs(tids []uint32, stats lids.Counter, m if order.IsReverse() { getBlockIndex = func(tid uint32) uint32 { return ti.lidsTable.GetLastBlockIndexForTID(tid) } getLIDsIterator = func(startIndex uint32, tid uint32) node.Node { - return (*lids.IteratorAsc)(lids.NewLIDsCursor(ti.lidsTable, ti.lidsLoader, startIndex, tid, stats, minLID, maxLID)) + return lids.NewIteratorAsc(ti.lidsTable, ti.lidsLoader, startIndex, tid, stats, minLID, maxLID) } } else { getBlockIndex = func(tid uint32) uint32 { return ti.lidsTable.GetFirstBlockIndexForTID(tid) } getLIDsIterator = func(startIndex uint32, tid uint32) node.Node { - return (*lids.IteratorDesc)(lids.NewLIDsCursor(ti.lidsTable, ti.lidsLoader, startIndex, tid, stats, minLID, maxLID)) + return lids.NewIteratorDesc(ti.lidsTable, ti.lidsLoader, startIndex, tid, stats, minLID, maxLID) } } diff --git a/frac/sealed_source.go b/frac/sealed_source.go index d537c16f..92c371f2 100644 --- a/frac/sealed_source.go +++ b/frac/sealed_source.go @@ -124,7 +124,7 @@ func (s *SealedSource) postingsForField(field string) iter.Seq2[indexwriter.Toke } chunkIdx := lidsTable.GetChunkIndex(bi, tid) - lidsBuf = append(lidsBuf, lidBlock.GetLIDs(chunkIdx)...) + lidsBuf = lidBlock.CopyLIDs(chunkIdx, lidsBuf) } if !yield(indexwriter.TokenLIDs{First: tokenVal, Second: lidsBuf}, nil) { diff --git a/fracmanager/config.go b/fracmanager/config.go index 6bc8acf2..7e6392d5 100644 --- a/fracmanager/config.go +++ b/fracmanager/config.go @@ -71,6 +71,9 @@ func FillConfigWithDefault(config *Config) *Config { if config.SealParams.TokenTableZstdLevel == 0 { config.SealParams.TokenTableZstdLevel = zstdDefaultLevel } + if config.SealParams.LIDsBitmapThreshold == 0 { + config.SealParams.LIDsBitmapThreshold = consts.DefaultLIDBlockCap + } if config.ReplayWorkers == 0 { config.ReplayWorkers = consts.DefaultReplayWorkers } diff --git a/indexwriter/blocks.go b/indexwriter/blocks.go index 7b2d6c0c..37bac25f 100644 --- a/indexwriter/blocks.go +++ b/indexwriter/blocks.go @@ -34,8 +34,8 @@ type lidExt struct { // unpackedLIDBlock represents a sealed block containing LID (Local ID) data. type unpackedLIDBlock struct { - ext lidExt // LIDs block metadata for registry marking - payload lids.Block // LID data payload + ext lidExt // LIDs block metadata for registry marking + payload lids.UnpackedBlock // LID data payload } // unpackedIDBlock represents a sealed block containing various identifier types. diff --git a/indexwriter/blocks_test.go b/indexwriter/blocks_test.go index 546d08a7..b567b84b 100644 --- a/indexwriter/blocks_test.go +++ b/indexwriter/blocks_test.go @@ -241,27 +241,27 @@ func TestBlocksBuilder_BuildTokenBlocks(t *testing.T) { expectedLIDBlocks := []unpackedLIDBlock{ { ext: lidExt{minTID: 1, maxTID: 1, isContinued: false}, - payload: lids.Block{LIDs: []uint32{10, 20, 30}, Offsets: []uint32{0, 3}, IsLastLID: false}, + payload: lids.UnpackedBlock{LIDs: []uint32{10, 20, 30}, Offsets: []uint32{0, 3}, IsLastLID: false}, }, { ext: lidExt{minTID: 1, maxTID: 3, isContinued: true}, - payload: lids.Block{LIDs: []uint32{40, 2, 3}, Offsets: []uint32{0, 1, 2, 3}, IsLastLID: true}, + payload: lids.UnpackedBlock{LIDs: []uint32{40, 2, 3}, Offsets: []uint32{0, 1, 2, 3}, IsLastLID: true}, }, { ext: lidExt{minTID: 4, maxTID: 6, isContinued: false}, - payload: lids.Block{LIDs: []uint32{4, 5, 6}, Offsets: []uint32{0, 1, 2, 3}, IsLastLID: true}, + payload: lids.UnpackedBlock{LIDs: []uint32{4, 5, 6}, Offsets: []uint32{0, 1, 2, 3}, IsLastLID: true}, }, { ext: lidExt{minTID: 7, maxTID: 9, isContinued: false}, - payload: lids.Block{LIDs: []uint32{7, 8, 9}, Offsets: []uint32{0, 1, 2, 3}, IsLastLID: true}, + payload: lids.UnpackedBlock{LIDs: []uint32{7, 8, 9}, Offsets: []uint32{0, 1, 2, 3}, IsLastLID: true}, }, { ext: lidExt{minTID: 10, maxTID: 12, isContinued: false}, - payload: lids.Block{LIDs: []uint32{10, 11, 12}, Offsets: []uint32{0, 1, 2, 3}, IsLastLID: true}, + payload: lids.UnpackedBlock{LIDs: []uint32{10, 11, 12}, Offsets: []uint32{0, 1, 2, 3}, IsLastLID: true}, }, { ext: lidExt{minTID: 13, maxTID: 14, isContinued: false}, - payload: lids.Block{LIDs: []uint32{13, 14}, Offsets: []uint32{0, 1, 2}, IsLastLID: true}, + payload: lids.UnpackedBlock{LIDs: []uint32{13, 14}, Offsets: []uint32{0, 1, 2}, IsLastLID: true}, }, } assert.Equal(t, expectedLIDBlocks, lidBlocks) diff --git a/indexwriter/index.go b/indexwriter/index.go index cdf6d1c5..424671cb 100644 --- a/indexwriter/index.go +++ b/indexwriter/index.go @@ -56,10 +56,10 @@ func (i indexBlock) bin(pos int64) (storage.IndexBlockHeader, []byte) { type IndexWriter struct { params common.SealParams - buf1 []byte - buf2 []byte - buf32 []uint32 - buf64 []uint64 + buf1 []byte + buf2 []byte + buf64 []uint64 + lidsPacker *lids.BlockPacker idsTable seqids.Table lidsTable lids.Table @@ -67,12 +67,17 @@ type IndexWriter struct { } func New(params common.SealParams) *IndexWriter { + lidsPacker := lids.NewBlockPacker() + if params.LIDsBitmapThreshold != 0 { + lidsPacker.LidsBitmapThreshold = params.LIDsBitmapThreshold + } + return &IndexWriter{ - params: params, - buf1: make([]byte, 0, consts.RegularBlockSize), - buf2: make([]byte, 0, consts.RegularBlockSize), - buf32: make([]uint32, 0, consts.DefaultLIDBlockCap), - buf64: make([]uint64, 0, consts.RegularBlockSize), + params: params, + buf1: make([]byte, 0, consts.RegularBlockSize), + buf2: make([]byte, 0, consts.RegularBlockSize), + buf64: make([]uint64, 0, consts.RegularBlockSize), + lidsPacker: lidsPacker, } } @@ -302,7 +307,7 @@ func (s *IndexWriter) packLIDsBlock(block unpackedLIDBlock) indexBlock { s.lidsTable.IsContinued = append(s.lidsTable.IsContinued, block.ext.isContinued) // Packing block - s.buf1 = block.payload.Pack(s.buf1[:0], s.buf32[:0]) + s.buf1 = s.lidsPacker.Pack(&block.payload, s.buf1[:0]) b := s.newIndexBlockZSTD(s.buf1, s.params.LIDsZstdLevel) b.ext1 = ext1 // Legacy continuation flag b.ext2 = uint64(block.ext.maxTID)<<32 | uint64(block.ext.minTID) // TID range diff --git a/indexwriter/lid_accumulator.go b/indexwriter/lid_accumulator.go index d2b1b97a..1c7c0a4d 100644 --- a/indexwriter/lid_accumulator.go +++ b/indexwriter/lid_accumulator.go @@ -5,6 +5,8 @@ import ( "github.com/ozontech/seq-db/frac/sealed/lids" ) +// lidAccumulator accumulates LIDs into blocks of fixed capacity. +// It implements the add function that receives []uint32 directly. type lidAccumulator struct { blockCapacity int onBlock func(unpackedLIDBlock) error @@ -30,7 +32,7 @@ func newLIDAccumulator( } a.currentBlock.ext.minTID = 1 - a.currentBlock.payload = lids.Block{ + a.currentBlock.payload = lids.UnpackedBlock{ LIDs: make([]uint32, 0, blockCapacity), Offsets: []uint32{0}, } diff --git a/node/batch.go b/node/batch.go index 92c8aa62..6da8ef8f 100644 --- a/node/batch.go +++ b/node/batch.go @@ -1,54 +1,326 @@ package node import ( - "slices" + "math" + "sort" + + "github.com/RoaringBitmap/roaring/v2" ) -// LIDBatch represents a batch of LIDs. lids are stored as uint32 slice and sorted in ascending order regardless of doc order. -// This allows to avoid copying and use reference to LID blocks data. -// Such batches are also logically immutable - we can't append or delete from them, only union or intersect. But we can zero out (reset) them. -type LIDBatch struct { - lids []uint32 - desc bool // if doc order is DESC (default order) +// LIDBatch is batch of lids. It's immutable and can not be modified. Lids are always +// sorted in ascending way for every underlying implementation. +type LIDBatch interface { + Len() int + IsEmpty() bool + // Min returns minimum (first) value. Panics if batch is empty. + Min() uint32 + // Max returns max (last) value. Panics if batch is empty. + Max() uint32 + CopyLIDs(desc bool, dst []LID) []LID + // Iter iterates lids in ascending way. + Iter() Iter + // ReverseIter iterates lids in descending way. + ReverseIter() Iter + // Narrow returns a batch containing only LIDs from minLID to maxLID (inclusive both). + Narrow(minLID, maxLID uint32) LIDBatch +} + +type Iter interface { + Next() (uint32, bool) + NextGeq(geq uint32) (uint32, bool) } -// NewDescBatch creates a batch of lids for DESC docs order -func NewDescBatch(lids []uint32) LIDBatch { - return LIDBatch{ - lids: lids, - desc: true, +func NewBitmapBatch(b *roaring.Bitmap) LIDBatch { + if b == nil || b.IsEmpty() { + return EmptyBatch() + } + return &bitmapBatch{ + bm: b, + min: b.Minimum(), + max: b.Maximum(), } } -// NewAscBatch creates a batch of lids for ASC docs order -func NewAscBatch(lids []uint32) LIDBatch { - return LIDBatch{ - lids: lids, - desc: false, +func NewBitmapBatchFromLids(lids []uint32) LIDBatch { + if len(lids) == 0 { + return EmptyBatch() } + b := roaring.NewBitmap() + b.AddMany(lids) + b.RunOptimize() + return NewBitmapBatch(b) } -func (b LIDBatch) Len() int { +func NewSliceBatch(lids []uint32) LIDBatch { + if len(lids) == 0 { + return EmptyBatch() + } + return &sliceBatch{lids: lids} +} + +// sliceBatch a batch of LIDs based on slice. LIDs are always sorted in ascending way. +// It's never empty. +type sliceBatch struct { + lids []uint32 +} + +func (b *sliceBatch) Len() int { return len(b.lids) } -func (b LIDBatch) LIDs(out []LID) []LID { - if b.desc { +func (b *sliceBatch) IsEmpty() bool { + return false +} + +func (b *sliceBatch) Min() uint32 { + return b.lids[0] +} + +func (b *sliceBatch) Max() uint32 { + return b.lids[len(b.lids)-1] +} + +func (b *sliceBatch) Narrow(minLID, maxLID uint32) LIDBatch { + batchMin := b.lids[0] + batchMax := b.lids[len(b.lids)-1] + if minLID <= batchMin && batchMax <= maxLID { + return b + } + if maxLID < batchMin || minLID > batchMax { + return EmptyBatch() + } + lo := 0 + if minLID > batchMin { + lo = sort.Search(len(b.lids), func(i int) bool { return b.lids[i] >= minLID }) + } + hi := len(b.lids) + if maxLID < batchMax { + hi = sort.Search(len(b.lids), func(i int) bool { return b.lids[i] > maxLID }) + } + if lo >= hi { + return EmptyBatch() + } + return &sliceBatch{lids: b.lids[lo:hi]} +} + +func (b *sliceBatch) Iter() Iter { + return &sliceIter{lids: b.lids, max: b.Max()} +} + +func (b *sliceBatch) ReverseIter() Iter { + return &sliceReverseIter{lids: b.lids, idx: len(b.lids) - 1} +} + +func (b *sliceBatch) CopyLIDs(desc bool, dst []LID) []LID { + if desc { for _, lid := range b.lids { - out = append(out, NewDescLID(lid)) + dst = append(dst, NewDescLID(lid)) + } + } else { + for i := len(b.lids) - 1; i >= 0; i-- { + dst = append(dst, NewAscLID(b.lids[i])) + } + } + return dst +} + +type sliceIter struct { + lids []uint32 + idx int + max uint32 +} + +func (it *sliceIter) Next() (uint32, bool) { + if it.idx >= len(it.lids) { + return 0, false + } + v := it.lids[it.idx] + it.idx++ + return v, true +} + +func (it *sliceIter) NextGeq(geq uint32) (uint32, bool) { + if it.idx >= len(it.lids) || (geq > it.max) { + it.idx = len(it.lids) + return 0, false + } + rest := it.lids[it.idx:] + off := sort.Search(len(rest), func(i int) bool { return rest[i] >= geq }) + it.idx += off + return it.Next() +} + +type sliceReverseIter struct { + lids []uint32 + idx int +} + +func (it *sliceReverseIter) Next() (uint32, bool) { + if it.idx < 0 { + return 0, false + } + v := it.lids[it.idx] + it.idx-- + return v, true +} + +func (it *sliceReverseIter) NextGeq(leq uint32) (uint32, bool) { + if it.idx < 0 { + return 0, false + } + right := it.idx + 1 + idx := sort.Search(right, func(i int) bool { return it.lids[i] > leq }) - 1 + if idx < 0 { + it.idx = -1 + return 0, false + } + it.idx = idx + return it.Next() +} + +// bitmapBatch a LIDs batch based on roaring bitmap. Never empty. +type bitmapBatch struct { + bm *roaring.Bitmap + min uint32 + max uint32 +} + +func (b *bitmapBatch) Len() int { + return int(b.bm.GetCardinality()) +} + +func (b *bitmapBatch) IsEmpty() bool { + return false +} + +func (b *bitmapBatch) Min() uint32 { + return b.min +} + +func (b *bitmapBatch) Max() uint32 { + return b.max +} + +func (b *bitmapBatch) Narrow(minLID, maxLID uint32) LIDBatch { + if minLID <= b.min && b.max <= maxLID { + return b + } + if maxLID < b.min || minLID > b.max { + return EmptyBatch() + } + // TODO(cheb0) use copy-on-write for bitmap? + out := b.bm.Clone() + if minLID > b.min { + out.RemoveRange(0, uint64(minLID)) + } + if maxLID < b.max { + out.RemoveRange(uint64(maxLID)+1, uint64(0x100000000)) + } + return NewBitmapBatch(out) +} + +func (b *bitmapBatch) Iter() Iter { + return newBitmapIter(b.bm) +} + +func (b *bitmapBatch) ReverseIter() Iter { + return newBitmapReverseIter(b.bm) +} + +func (b *bitmapBatch) CopyLIDs(desc bool, dst []LID) []LID { + if desc { + it := b.bm.Iterator() + for it.HasNext() { + dst = append(dst, NewDescLID(it.Next())) } } else { - for _, lid := range slices.Backward(b.lids) { - out = append(out, NewAscLID(lid)) + it := b.bm.ReverseIterator() + for it.HasNext() { + dst = append(dst, NewAscLID(it.Next())) } } + return dst +} - return out +type emptyBatch struct{} + +var emptyBatchInstance = emptyBatch{} + +func EmptyBatch() LIDBatch { + return emptyBatchInstance +} + +func (emptyBatch) Len() int { return 0 } +func (emptyBatch) IsEmpty() bool { return true } + +func (emptyBatch) Min() uint32 { + panic("min called on empty batch") +} + +func (emptyBatch) Max() uint32 { + panic("Maximum called on empty batch") +} + +func (emptyBatch) Narrow(uint32, uint32) LIDBatch { return emptyBatchInstance } +func (emptyBatch) CopyLIDs(_ bool, dst []LID) []LID { return dst } +func (emptyBatch) Iter() Iter { return emptyIterInstance } +func (emptyBatch) ReverseIter() Iter { return emptyIterInstance } + +type emptyIter struct{} + +var emptyIterInstance = emptyIter{} + +func (emptyIter) Next() (uint32, bool) { return 0, false } +func (emptyIter) NextGeq(uint32) (uint32, bool) { return 0, false } + +type bitmapIter struct { + it roaring.IntIterator +} + +func newBitmapIter(b *roaring.Bitmap) *bitmapIter { + var it roaring.IntIterator + it.Initialize(b) + return &bitmapIter{it: it} +} + +func (f *bitmapIter) Next() (uint32, bool) { + if !f.it.HasNext() { + return 0, false + } + return f.it.Next(), true +} + +func (f *bitmapIter) NextGeq(geq uint32) (uint32, bool) { + f.it.AdvanceIfNeeded(geq) + if !f.it.HasNext() { + return 0, false + } + return f.it.Next(), true +} + +type bitmapReverseIter struct { + bm *roaring.Bitmap + pos uint32 +} + +func newBitmapReverseIter(bm *roaring.Bitmap) *bitmapReverseIter { + return &bitmapReverseIter{bm: bm, pos: math.MaxUint32} +} + +func (it *bitmapReverseIter) Next() (uint32, bool) { + prev := it.bm.PreviousValue(it.pos - 1) + if prev == -1 { + return 0, false + } + it.pos = uint32(prev) + return uint32(prev), true } -func (b LIDBatch) Reset() LIDBatch { - return LIDBatch{ - lids: b.lids[:0], - desc: b.desc, +func (it *bitmapReverseIter) NextGeq(leq uint32) (uint32, bool) { + prev := it.bm.PreviousValue(min(it.pos-1, leq)) + if prev == -1 { + return 0, false } + it.pos = uint32(prev) + return uint32(prev), true } diff --git a/node/batch_test.go b/node/batch_test.go new file mode 100644 index 00000000..1454b538 --- /dev/null +++ b/node/batch_test.go @@ -0,0 +1,242 @@ +package node + +import ( + "math" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type batchFactory func([]uint32) LIDBatch + +var batchFactories = []struct { + name string + build batchFactory +}{ + {name: "bitmap", build: NewBitmapBatchFromLids}, + {name: "slice", build: NewSliceBatch}, +} + +func TestLIDBatchNarrow(t *testing.T) { + testCases := []struct { + name string + input []uint32 + minLID uint32 + maxLID uint32 + expected []uint32 + }{ + { + name: "empty batch", + input: nil, + minLID: 0, + maxLID: 10, + expected: nil}, + { + name: "full range no-op", + input: []uint32{1, 5, 10}, + minLID: 0, + maxLID: math.MaxUint32, + expected: []uint32{1, 5, 10}}, + { + name: "trim below", + input: []uint32{1, 5, 10, 15}, + minLID: 6, + maxLID: 20, + expected: []uint32{10, 15}}, + { + name: "exact min boundary", + input: []uint32{1, 5, 10, 15}, + minLID: 10, + maxLID: 20, + expected: []uint32{10, 15}}, + { + name: "trim above", + input: []uint32{1, 5, 10, 15}, + minLID: 0, + maxLID: 10, + expected: []uint32{1, 5, 10}}, + { + name: "trim both sides", + input: []uint32{1, 5, 10, 15, 20}, + minLID: 5, + maxLID: 15, + expected: []uint32{5, 10, 15}}, + { + name: "no overlap below", + input: []uint32{1, 5, 10}, + minLID: 11, + maxLID: 20, + expected: nil}, + { + name: "no overlap above", + input: []uint32{5, 10, 20}, + minLID: 0, + maxLID: 4, + expected: nil}, + { + name: "overlap at min", + input: []uint32{7, 10, 20, 25}, + minLID: 0, + maxLID: 7, + expected: []uint32{7}}, + { + name: "overlap at max", + input: []uint32{7, 10, 20, 25}, + minLID: 25, + maxLID: 40, + expected: []uint32{25}}, + } + + for _, impl := range batchFactories { + t.Run(impl.name, func(t *testing.T) { + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + b := impl.build(tc.input) + got := b.Narrow(tc.minLID, tc.maxLID) + assert.Equal(t, tc.expected, toSlice(got)) + }) + } + }) + } +} + +func TestBitmapBatchNarrow_NoCloneWhenUnchanged(t *testing.T) { + src := NewBitmapBatchFromLids([]uint32{1, 5, 10, 15, 20}) + got := src.Narrow(1, 20) + assert.Same(t, src, got) +} + +func toSlice(b LIDBatch) []uint32 { + if b == nil || b.IsEmpty() { + return nil + } + it := b.Iter() + out := make([]uint32, 0, b.Len()) + for { + v, ok := it.Next() + if !ok { + break + } + out = append(out, v) + } + return out +} + +func TestNextGeq(t *testing.T) { + for _, impl := range batchFactories { + t.Run(impl.name, func(t *testing.T) { + b := impl.build([]uint32{1, 5, 10, 15, 20, 21, 22, 26, 30}) + it := b.Iter() + + v, ok := it.NextGeq(1) + require.True(t, ok) + assert.Equal(t, uint32(1), v) + + // calling NextGeq with already seen value returns next value + v, ok = it.NextGeq(1) + require.True(t, ok) + assert.Equal(t, uint32(5), v) + + v, ok = it.NextGeq(13) + require.True(t, ok) + assert.Equal(t, uint32(15), v) + + v, ok = it.NextGeq(22) + require.True(t, ok) + assert.Equal(t, uint32(22), v) + + _, ok = it.NextGeq(50) + assert.False(t, ok) + }) + } +} + +func TestReverseNextGeq(t *testing.T) { + for _, impl := range batchFactories { + t.Run(impl.name, func(t *testing.T) { + b := impl.build([]uint32{3, 5, 10, 15, 20, 21, 22, 26, 30}) + it := b.ReverseIter() + + v, ok := it.NextGeq(1000) + require.True(t, ok) + assert.Equal(t, uint32(30), v) + + // calling NextGeq with already seen value returns next value + v, ok = it.NextGeq(30) + require.True(t, ok) + assert.Equal(t, uint32(26), v) + + v, ok = it.NextGeq(20) + require.True(t, ok) + assert.Equal(t, uint32(20), v) + + v, ok = it.NextGeq(9) + require.True(t, ok) + assert.Equal(t, uint32(5), v) + + _, ok = it.NextGeq(2) + assert.False(t, ok) + }) + } +} + +func TestBatchIter(t *testing.T) { + for _, impl := range batchFactories { + t.Run(impl.name, func(t *testing.T) { + b := impl.build([]uint32{1, 5, 10, 15, 20}) + it := b.Iter() + + var got []uint32 + for { + v, ok := it.Next() + if !ok { + break + } + got = append(got, v) + } + assert.Equal(t, []uint32{1, 5, 10, 15, 20}, got) + + b = impl.build([]uint32{1, 5, 10, 15, 20}) + it = b.Iter() + v, ok := it.NextGeq(11) + require.True(t, ok) + assert.Equal(t, uint32(15), v) + + b = impl.build([]uint32{1, 5, 10}) + it = b.Iter() + _, ok = it.NextGeq(100) + assert.False(t, ok) + }) + } +} + +func TestBatchReverseIter(t *testing.T) { + for _, impl := range batchFactories { + t.Run(impl.name, func(t *testing.T) { + b := impl.build([]uint32{1, 5, 10, 15, 20}) + it := b.ReverseIter() + + var got []uint32 + for { + v, ok := it.Next() + if !ok { + break + } + got = append(got, v) + } + assert.Equal(t, []uint32{20, 15, 10, 5, 1}, got) + + b = impl.build([]uint32{1, 5, 10, 15, 20}) + it = b.ReverseIter() + v, ok := it.NextGeq(11) + require.True(t, ok) + assert.Equal(t, uint32(10), v) + + b = impl.build([]uint32{1, 5, 10}) + it = b.ReverseIter() + _, ok = it.NextGeq(0) + assert.False(t, ok) + }) + } +} diff --git a/node/node.go b/node/node.go index 5334f218..98cf21e3 100644 --- a/node/node.go +++ b/node/node.go @@ -14,9 +14,9 @@ type Node interface { type BatchedNode interface { fmt.Stringer // NextBatch returns next batch. Returns nil when exhausted. - NextBatch() LIDBatch + NextBatch(need int) LIDBatch // NextBatchGeq returns next batch (LIDs >= minLID). Returns nil when exhausted. - NextBatchGeq(nextLID LID) LIDBatch + NextBatchGeq(need int, nextLID LID) LIDBatch } type Sourced interface {