From c42261ca1770a31faf336275d0d0298c431dcefe Mon Sep 17 00:00:00 2001 From: Andrei Cheboksarov <37665782+cheb0@users.noreply.github.com> Date: Thu, 30 Jul 2026 11:46:12 +0400 Subject: [PATCH] enhance tokens block memory usage --- config/frac_version.go | 5 +- frac/sealed/token/block_loader.go | 83 ++++++++++++++++++++++++-- frac/sealed/token/block_loader_test.go | 52 ++++++---------- indexwriter/blocks.go | 12 ++-- 4 files changed, 108 insertions(+), 44 deletions(-) diff --git a/config/frac_version.go b/config/frac_version.go index d6bf41d6..35260124 100644 --- a/config/frac_version.go +++ b/config/frac_version.go @@ -24,6 +24,9 @@ const ( // BinaryDataV5 - token blocks have zone maps (eng letters presense) and doc frequencies for heavy tokens BinaryDataV5 + + // BinaryDataV6 - keep offsets separate in token block + BinaryDataV6 ) -const CurrentFracVersion = BinaryDataV5 +const CurrentFracVersion = BinaryDataV6 diff --git a/frac/sealed/token/block_loader.go b/frac/sealed/token/block_loader.go index 26631572..8fd444e2 100644 --- a/frac/sealed/token/block_loader.go +++ b/frac/sealed/token/block_loader.go @@ -24,6 +24,9 @@ type Block struct { Offsets []uint32 FreqIndexes []uint16 // indexes of tokens which have doc freqs (frequencies) Freqs []uint32 // frequencies of certain tokens (how many docs have this token included at least once) + + // TODO(cheb0) delete this field and convert V0..V5 to a new in-memory format when data all clusters have V6 fractions + FracVer config.BinaryDataVersion } func (b *Block) Size() int { @@ -46,6 +49,8 @@ func (b Block) Pack(dst []byte, buf []uint32) []byte { dst = binary.LittleEndian.AppendUint32(dst, uint32(len(b.Payload))) dst = append(dst, b.Payload...) + dst = packer.CompressDeltaBitpackUint32(dst, b.Offsets, buf) + if len(b.FreqIndexes) > 0 { dst = packer.CompressDeltaBitpackUint16(dst, b.FreqIndexes, buf) dst = packer.CompressDeltaBitpackUint32(dst, b.Freqs, buf) @@ -55,6 +60,13 @@ func (b Block) Pack(dst []byte, buf []uint32) []byte { } func (b *Block) Unpack(data []byte, fracVer config.BinaryDataVersion, unpackBuf *UnpackBuffer) error { + b.FracVer = fracVer + + if fracVer >= config.BinaryDataV6 { + unpackBuf.Reset(fracVer) + return b.unpackV6(data, unpackBuf) + } + if fracVer >= config.BinaryDataV5 { unpackBuf.Reset(fracVer) return b.unpackV5(data, unpackBuf) @@ -62,9 +74,39 @@ func (b *Block) Unpack(data []byte, fracVer config.BinaryDataVersion, unpackBuf return b.unpackV1(data) } -func (b *Block) unpackV1(data []byte) error { - b.Payload = append([]byte{}, data...) - return b.parseTokenPayload(b.Payload) +func (b *Block) unpackV6(data []byte, buf *UnpackBuffer) error { + if len(data) < util.SizeOfUint32 { + return fmt.Errorf("token block too short: %d bytes", len(data)) + } + flags := data[0] + data = data[1:] + + // token payload + payloadLen := binary.LittleEndian.Uint32(data[:util.SizeOfUint32]) + data = data[util.SizeOfUint32:] + if uint32(len(data)) < payloadLen { + return fmt.Errorf("invalid token block payload length: %d, data len %d", payloadLen, len(data)) + } + + payload := data[:payloadLen] + data = data[payloadLen:] + + b.Payload = append(b.Payload[:0], payload...) + + // offsets + var err error + data, buf.decompressedUint32, err = packer.DecompressDeltaBitpackUint32(data, buf.decompressedUint32, buf.compressed) + if err != nil { + return err + } + b.Offsets = append(b.Offsets, buf.decompressedUint32...) + + err = b.unpackFreqs(data, buf, flags) + if err != nil { + return err + } + + return nil } func (b *Block) unpackV5(data []byte, buf *UnpackBuffer) error { @@ -85,11 +127,22 @@ func (b *Block) unpackV5(data []byte, buf *UnpackBuffer) error { b.Payload = append(b.Payload[:0], payload...) - if err := b.parseTokenPayload(payload); err != nil { + if err := b.parseTokenPayloadV5(payload); err != nil { return err } + err := b.unpackFreqs(data, buf, flags) + if err != nil { + return err + } + + return nil +} + +func (b *Block) unpackFreqs(data []byte, buf *UnpackBuffer, flags byte) error { if flags&1 > 0 { + buf.decompressedUint32 = buf.decompressedUint32[:0] + var err error data, buf.decompressedUint16, err = packer.DecompressDeltaBitpackUint16(data, buf.decompressedUint16, buf.compressed) if err != nil { @@ -103,11 +156,16 @@ func (b *Block) unpackV5(data []byte, buf *UnpackBuffer) error { } b.Freqs = append(b.Freqs, buf.decompressedUint32...) } - return nil } -func (b *Block) parseTokenPayload(data []byte) error { +func (b *Block) unpackV1(data []byte) error { + b.Payload = append([]byte{}, data...) + return b.parseTokenPayloadV5(b.Payload) +} + +// parseTokenPayloadV5 derives offsets from tokens payload. Only used for v1...v5 legacy fractions. +func (b *Block) parseTokenPayloadV5(data []byte) error { b.Offsets = b.Offsets[:0] var offset uint32 @@ -129,6 +187,10 @@ func (b *Block) parseTokenPayload(data []byte) error { } func (b *Block) Len() int { + if b.FracVer >= config.BinaryDataV6 { + return len(b.Offsets) - 1 + } + return len(b.Offsets) } @@ -147,6 +209,15 @@ func (b *Block) GetFreq(index int) uint32 { } func (b *Block) GetToken(index int) []byte { + if b.FracVer >= config.BinaryDataV6 { + return b.Payload[b.Offsets[index]:b.Offsets[index+1]] + } + + return b.getTokenV5(index) +} + +//go:noinline +func (b *Block) getTokenV5(index int) []byte { offset := b.Offsets[index] l := binary.LittleEndian.Uint32(b.Payload[offset:]) offset += uint32(util.SizeOfUint32) // skip val length diff --git a/frac/sealed/token/block_loader_test.go b/frac/sealed/token/block_loader_test.go index 6f45e02a..1e6b052f 100644 --- a/frac/sealed/token/block_loader_test.go +++ b/frac/sealed/token/block_loader_test.go @@ -1,7 +1,6 @@ package token import ( - "encoding/binary" "testing" "github.com/stretchr/testify/assert" @@ -11,14 +10,12 @@ import ( ) func TestBlock_PackUnpack_NoFreq(t *testing.T) { - src := Block{ - Payload: packTokenPayload([]byte("foo"), []byte("bar")), - } + src := buildTokenPayload([]byte("foo"), []byte("bar")) var buf []uint32 packed := src.Pack(nil, buf) var dst Block - require.NoError(t, dst.Unpack(packed, config.BinaryDataV5, &UnpackBuffer{})) + require.NoError(t, dst.Unpack(packed, config.BinaryDataV6, &UnpackBuffer{})) assert.Equal(t, 2, dst.Len()) assert.Equal(t, []byte("foo"), dst.GetToken(0)) @@ -29,16 +26,14 @@ func TestBlock_PackUnpack_NoFreq(t *testing.T) { } func TestBlock_PackUnpack_WithFreq(t *testing.T) { - src := Block{ - Payload: packTokenPayload([]byte("dog"), []byte("cat"), []byte("horse"), []byte("duck")), - FreqIndexes: []uint16{0, 2}, - Freqs: []uint32{100, 200}, - } + src := buildTokenPayload([]byte("dog"), []byte("cat"), []byte("horse"), []byte("duck")) + src.FreqIndexes = []uint16{0, 2} + src.Freqs = []uint32{100, 200} var buf []uint32 packed := src.Pack(nil, buf) var dst Block - require.NoError(t, dst.Unpack(packed, config.BinaryDataV5, &UnpackBuffer{})) + require.NoError(t, dst.Unpack(packed, config.BinaryDataV6, &UnpackBuffer{})) assert.Equal(t, src.Payload, dst.Payload) @@ -48,24 +43,10 @@ func TestBlock_PackUnpack_WithFreq(t *testing.T) { assert.Equal(t, uint32(0), dst.GetFreq(3)) } -func TestBlock_Unpack_Legacy(t *testing.T) { - legacy := packTokenPayload([]byte("legacy")) - - var dst Block - require.NoError(t, dst.Unpack(legacy, config.BinaryDataV4, &UnpackBuffer{})) - - assert.Equal(t, legacy, dst.Payload) - assert.Equal(t, []uint32{0}, dst.Offsets) - assert.Empty(t, dst.FreqIndexes) - assert.Empty(t, dst.Freqs) -} - func TestBlock_UnpackBufferReuse(t *testing.T) { - src := Block{ - Payload: packTokenPayload([]byte("a"), []byte("b")), - FreqIndexes: []uint16{1}, - Freqs: []uint32{64}, - } + src := buildTokenPayload([]byte("a"), []byte("b")) + src.FreqIndexes = []uint16{1} + src.Freqs = []uint32{64} var packBuf []uint32 packed := src.Pack(nil, packBuf) @@ -73,8 +54,8 @@ func TestBlock_UnpackBufferReuse(t *testing.T) { unpackBuf := &UnpackBuffer{} var dst1, dst2 Block - require.NoError(t, dst1.Unpack(packed, config.BinaryDataV5, unpackBuf)) - require.NoError(t, dst2.Unpack(packed, config.BinaryDataV5, unpackBuf)) + require.NoError(t, dst1.Unpack(packed, config.BinaryDataV6, unpackBuf)) + require.NoError(t, dst2.Unpack(packed, config.BinaryDataV6, unpackBuf)) assert.Equal(t, dst1.FreqIndexes, dst2.FreqIndexes) assert.Equal(t, dst1.Freqs, dst2.Freqs) @@ -83,11 +64,16 @@ func TestBlock_UnpackBufferReuse(t *testing.T) { assert.Equal(t, uint32(64), dst2.GetFreq(1)) } -func packTokenPayload(tokens ...[]byte) []byte { +func buildTokenPayload(tokens ...[]byte) Block { var payload []byte + var offsets []uint32 + offsets = append(offsets, 0) for _, tok := range tokens { - payload = binary.LittleEndian.AppendUint32(payload, uint32(len(tok))) + offsets = append(offsets, offsets[len(offsets)-1]+uint32(len(tok))) payload = append(payload, tok...) } - return payload + return Block{ + Payload: payload, + Offsets: offsets, + } } diff --git a/indexwriter/blocks.go b/indexwriter/blocks.go index a05dc262..13b9c51a 100644 --- a/indexwriter/blocks.go +++ b/indexwriter/blocks.go @@ -1,11 +1,11 @@ package indexwriter import ( - "encoding/binary" "iter" "math" "unsafe" + "github.com/ozontech/seq-db/config" "github.com/ozontech/seq-db/frac/sealed/lids" "github.com/ozontech/seq-db/frac/sealed/seqids" "github.com/ozontech/seq-db/frac/sealed/token" @@ -56,6 +56,7 @@ func tokenBlock( blockIdx uint32 blockSize int ) + block.payload.FracVer = config.BinaryDataV6 var ( currentTID uint32 @@ -123,9 +124,12 @@ func tokenBlock( } } - tokenIndex := uint32(len(block.payload.Offsets)) - block.payload.Offsets = append(block.payload.Offsets, uint32(len(block.payload.Payload))) - block.payload.Payload = binary.LittleEndian.AppendUint32(block.payload.Payload, uint32(len(tok))) + offsets := block.payload.Offsets + if len(offsets) == 0 { + offsets = append(offsets, 0) + } + tokenIndex := uint32(len(offsets) - 1) + block.payload.Offsets = append(offsets, offsets[len(offsets)-1]+uint32(len(tok))) block.payload.Payload = append(block.payload.Payload, tok...) if len(tlids) >= tokenFreqAbsThreshold {