diff --git a/fragment.go b/fragment.go index 5b861a53d..94d65dc77 100644 --- a/fragment.go +++ b/fragment.go @@ -1731,6 +1731,42 @@ func (f *fragment) bulkImportMutex(rowIDs, columnIDs []uint64) error { return errors.Wrap(f.importPositions(toSet, toClear, rowSet), "importing positions") } +func (f *fragment) importValueSmallWrite(columnIDs, values []uint64, bitDepth uint, clear bool) error { + // TODO figure out how to avoid re-allocating these each time. Probably + // possible to store them on the fragment with a capacity based on + // MaxOpN. For now, we know that the total number of bits to be + // set+cleared is len(values)*(bitDepth+1), so we make each slice + // slightly more than half of that to try to avoid reallocation. + toSet := make([]uint64, 0, len(columnIDs)*int(bitDepth+1)*(5/8)) + toClear := make([]uint64, 0, len(columnIDs)*int(bitDepth+1)*(5/8)) + colSet := make(map[uint64]struct{}, len(columnIDs)) + + if err := func() (err error) { + for i := len(columnIDs) - 1; i >= 0; i-- { + columnID, value := columnIDs[i], values[i] + if _, ok := colSet[columnID]; ok { + continue + } + colSet[columnID] = struct{}{} + toSet, toClear, err = f.positionsForValue(columnID, bitDepth, value, clear, toSet, toClear) + if err != nil { + return errors.Wrap(err, "getting positions for value") + } + } + return nil + }(); err != nil { + _ = f.closeStorage() + _ = f.openStorage() + return err + } + rowSet := make(map[uint64]struct{}, bitDepth+1) + for i := uint(0); i < bitDepth+1; i++ { + rowSet[uint64(i)] = struct{}{} + } + err := f.importPositions(toSet, toClear, rowSet) + return errors.Wrap(err, "importing positions") +} + // importValue bulk imports a set of range-encoded values. func (f *fragment) importValue(columnIDs, values []uint64, bitDepth uint, clear bool) error { f.mu.Lock() @@ -1740,37 +1776,21 @@ func (f *fragment) importValue(columnIDs, values []uint64, bitDepth uint, clear if len(columnIDs) != len(values) { return fmt.Errorf("mismatch of column/value len: %d != %d", len(columnIDs), len(values)) } - var toSet, toClear []uint64 - smallWrite := false if len(columnIDs)*int(bitDepth+1)+f.opN < f.MaxOpN { - smallWrite = true + return errors.Wrap(f.importValueSmallWrite(columnIDs, values, bitDepth, clear), "import small write") - // TODO figure out how to avoid re-allocating these each time. Probably - // possible to store them on the fragment with a capacity based on - // MaxOpN. For now, we know that the total number of bits to be - // set+cleared is len(values)*(bitDepth+1), so we make each slice - // slightly more than half of that to try to avoid reallocation. - toSet = make([]uint64, 0, len(columnIDs)*int(bitDepth+1)*(5/8)) - toClear = make([]uint64, 0, len(columnIDs)*int(bitDepth+1)*(5/8)) } - if !smallWrite { - f.storage.OpWriter = nil - } + f.storage.OpWriter = nil // Process every value. // If an error occurs then reopen the storage. if err := func() (err error) { for i := range columnIDs { columnID, value := columnIDs[i], values[i] - if smallWrite { - toSet, toClear, err = f.positionsForValue(columnID, bitDepth, value, clear, toSet, toClear) - if err != nil { - return errors.Wrap(err, "getting positions for value") - } - } else if _, err = f.importSetValue(columnID, bitDepth, value, clear); err != nil { - return errors.Wrapf(err, "setting value, smallPath:%v", smallWrite) + if _, err := f.importSetValue(columnID, bitDepth, value, clear); err != nil { + return errors.Wrapf(err, "importSetValue") } } return nil @@ -1780,14 +1800,6 @@ func (f *fragment) importValue(columnIDs, values []uint64, bitDepth uint, clear return err } - if smallWrite { - rowSet := make(map[uint64]struct{}, bitDepth+1) - for i := uint(0); i < bitDepth+1; i++ { - rowSet[uint64(i)] = struct{}{} - } - err := f.importPositions(toSet, toClear, rowSet) - return errors.Wrap(err, "importing positions") - } err := f.snapshot() return errors.Wrap(err, "snapshotting") } diff --git a/fragment_internal_test.go b/fragment_internal_test.go index d553ca0fa..4df315f8a 100644 --- a/fragment_internal_test.go +++ b/fragment_internal_test.go @@ -3070,3 +3070,50 @@ func TestImportValueConcurrent(t *testing.T) { t.Fatalf("concurrently importing values: %v", err) } } + +func TestImportMultipleValues(t *testing.T) { + tests := []struct { + cols []uint64 + vals []uint64 + checkCols []uint64 + checkVals []uint64 + depth uint + }{ + { + cols: []uint64{0, 0}, + vals: []uint64{97, 100}, + depth: 7, + checkCols: []uint64{0}, + checkVals: []uint64{100}, + }, + } + + for i, test := range tests { + for _, maxOpN := range []int{0, 10000} { // test small/large write + t.Run(fmt.Sprintf("%dLowOpN", i), func(t *testing.T) { + f := mustOpenFragment("i", "f", viewBSIGroupPrefix+"foo", 0, CacheTypeNone) + f.MaxOpN = maxOpN + defer f.Clean(t) + err := f.importValue(test.cols, test.vals, test.depth, false) + if err != nil { + t.Fatalf("importing values: %v", err) + } + + for i := range test.checkCols { + cc, cv := test.checkCols[i], test.checkVals[i] + n, exists, err := f.value(cc, test.depth) + if err != nil { + t.Fatalf("getting value: %v", err) + } + if !exists { + t.Errorf("column %d should exist", cc) + } + if n != 100 { + t.Errorf("wrong value: %d is not %d", n, cv) + } + } + }) + + } + } +}