diff --git a/fragment.go b/fragment.go index 832d02f32..9922510ac 100644 --- a/fragment.go +++ b/fragment.go @@ -1058,7 +1058,7 @@ func (f *fragment) setValueBase(tx Tx, columnID uint64, bitDepth uint, value int } // importSetValue is a more efficient SetValue just for imports. -func (f *fragment) importSetValue(tx Tx, columnID uint64, bitDepth uint, value int64, clear bool) (changed int, err error) { // nolint: unparam +func (f *fragment) importSetValue(txb *TxBitmap, columnID uint64, bitDepth uint, value int64, clear bool) (changed int, err error) { // nolint: unparam // Convert value to an unsigned representation. uvalue := uint64(value) if value < 0 { @@ -1072,17 +1072,17 @@ func (f *fragment) importSetValue(tx Tx, columnID uint64, bitDepth uint, value i } if uvalue&(1< 0 { + } else if c { changed++ } } else { - changeCount, err := tx.Remove(f.index, f.field, f.view, f.shard, bit) + c, err := txb.Remove(bit) if err != nil { return changed, errors.Wrap(err, "removing") - } else if changeCount > 0 { + } else if c { changed++ } } @@ -1092,15 +1092,15 @@ func (f *fragment) importSetValue(tx Tx, columnID uint64, bitDepth uint, value i if p, err := f.pos(uint64(bsiExistsBit), columnID); err != nil { return changed, errors.Wrap(err, "getting not-null pos") } else if clear { - if c, err := tx.Remove(f.index, f.field, f.view, f.shard, p); err != nil { + if c, err := txb.Remove(p); err != nil { return changed, errors.Wrap(err, "removing not-null from storage") - } else if c > 0 { + } else if c { changed++ } } else { - if c, err := tx.Add(f.index, f.field, f.view, f.shard, !doBatched, p); err != nil { + if c, err := txb.Add(p); err != nil { return changed, errors.Wrap(err, "adding not-null to storage") - } else if c > 0 { + } else if c { changed++ } } @@ -1109,15 +1109,15 @@ func (f *fragment) importSetValue(tx Tx, columnID uint64, bitDepth uint, value i if p, err := f.pos(uint64(bsiSignBit), columnID); err != nil { return changed, errors.Wrap(err, "getting sign pos") } else if value >= 0 || clear { - if c, err := tx.Remove(f.index, f.field, f.view, f.shard, p); err != nil { + if c, err := txb.Remove(p); err != nil { return changed, errors.Wrap(err, "removing sign from storage") - } else if c > 0 { + } else if c { changed++ } } else { - if c, err := tx.Add(f.index, f.field, f.view, f.shard, !doBatched, p); err != nil { + if c, err := txb.Add(p); err != nil { return changed, errors.Wrap(err, "adding sign to storage") - } else if c > 0 { + } else if c { changed++ } } @@ -2460,17 +2460,20 @@ func (f *fragment) importValue(tx Tx, columnIDs []uint64, values []int64, bitDep if f.storage != nil { f.storage.OpWriter = nil } - totalChanges := 0 + + var totalChanges int if err := func() (err error) { + // Build changes into temporary bitmap. + txb := NewTxBitmap(tx, f.index, f.field, f.view, f.shard) for i := range columnIDs { columnID, value := columnIDs[i], values[i] - changed, err := f.importSetValue(tx, columnID, bitDepth, value, clear) - if err != nil { + if _, err := f.importSetValue(txb, columnID, bitDepth, value, clear); err != nil { return errors.Wrapf(err, "importSetValue") } - totalChanges += changed } - return nil + + // Flush changes in bulk back to the transaction. + return txb.Flush() }(); err != nil { _ = f.openStorage(true) return err diff --git a/tx.go b/tx.go index a5c154654..2b74b8fa6 100644 --- a/tx.go +++ b/tx.go @@ -231,3 +231,68 @@ type RawRoaringData struct { func (rr *RawRoaringData) Iterator() (roaring.RoaringIterator, error) { return roaring.NewRoaringIterator(rr.data) } + +// TxBitmap represents a bitmap that acts as a cache in front of a transaction. +// Updates to the bitmap first pull in containers as needed and update them +// in memory. The changes can be flushed in bulk using Flush(). +type TxBitmap struct { + b *roaring.Bitmap + tx Tx + index string + field string + view string + shard uint64 +} + +func NewTxBitmap(tx Tx, index, field, view string, shard uint64) *TxBitmap { + return &TxBitmap{ + b: roaring.NewBitmap(), + tx: tx, + index: index, + field: field, + view: view, + shard: shard, + } +} + +func (b *TxBitmap) Add(a ...uint64) (changed bool, err error) { + if err := b.ensureContainers(a...); err != nil { + return false, err + } + return b.b.Add(a...) +} + +func (b *TxBitmap) Remove(a ...uint64) (changed bool, err error) { + if err := b.ensureContainers(a...); err != nil { + return false, err + } + return b.b.Remove(a...) +} + +// ensureContainers pulls containers in from the transaction, if needed. +func (b *TxBitmap) ensureContainers(a ...uint64) error { + for _, v := range a { + key := highbits(v) + if b.b.Containers.Get(key) != nil { + continue + } + + c, err := b.tx.Container(b.index, b.field, b.view, b.shard, key) + if err != nil { + return err + } + b.b.Containers.Put(key, c) + } + return nil +} + +// Flush writes all containers in the bitmap back to the transaction. +func (b *TxBitmap) Flush() error { + for it, _ := b.b.Containers.Iterator(0); it.Next(); { + key, c := it.Value() + if err := b.tx.PutContainer(b.index, b.field, b.view, b.shard, key, c); err != nil { + return err + } + } + return nil +}