ensure mutex imports unset previous columns

fragment.go

@@ -1330,16 +1330,29 @@ func (f *fragment) bulkImport(rowIDs, columnIDs []uint64) error {
 		return fmt.Errorf("mismatch of row/column len: %d != %d", len(rowIDs), len(columnIDs))
 	}
 
+	if f.mutexVector != nil {
+		return f.bulkImportMutex(rowIDs, columnIDs)
+	}
+	return f.bulkImportStandard(rowIDs, columnIDs)
+}
+
+// bulkImportStandard performs a bulk import on a standard fragment.
+func (f *fragment) bulkImportStandard(rowIDs, columnIDs []uint64) error {
 	// Create a temporary bitmap which will be populated by rowIDs and columnIDs
 	// and then merged into the existing fragment's bitmap.
 	localBitmap := roaring.NewBitmap()
+
 	// Disconnect op writer so we don't append updates.
 	localBitmap.OpWriter = nil
 
-	// Process every bit.
-	// If an error occurs then reopen the storage.
-	lastID := uint64(0)
+	// rowSet maintains the set of rowIDs present in this import.
+	// It allows the cache to be updated once per row, instead of once
+	// per bit.
 	rowSet := make(map[uint64]struct{})
+	lastRowID := uint64(0)
+
+	// Process every bit by writing to a local bitmap,
+	// to be merged with fragment storage next.
 	for i := range rowIDs {
 		rowID, columnID := rowIDs[i], columnIDs[i]
 
@@ -1349,18 +1362,18 @@ func (f *fragment) bulkImport(rowIDs, columnIDs []uint64) error {
 			return err
 		}
 
-		// Write to storage.
+		// Write to local storage.
 		_, err = localBitmap.Add(pos)
 		if err != nil {
 			return err
 		}
+
 		// Reduce the StatsD rate for high volume stats
 		f.stats.Count("ImportBit", 1, 0.0001)
-		// import optimization to avoid linear foreach calls
-		// slight risk of concurrent cache counter being off but
-		// no real danger
-		if i == 0 || rowID != lastID {
-			lastID = rowID
+
+		// Add row to rowSet.
+		if i == 0 || rowID != lastRowID {
+			lastRowID = rowID
 			rowSet[rowID] = struct{}{}
 		}
 
@@ -1389,6 +1402,96 @@ func (f *fragment) bulkImport(rowIDs, columnIDs []uint64) error {
 	return unprotectedWriteToFragment(f, results)
 }
 
+// bulkImportMutex performs a bulk import on a fragment while ensuring
+// mutex restrictions. Because the mutex requirements must be checked
+// against storage, this method must acquire a write lock on the fragment
+// during the entire process, and it handles every bit independently.
+func (f *fragment) bulkImportMutex(rowIDs, columnIDs []uint64) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	// Disconnect op writer so we don't append updates.
+	f.storage.OpWriter = nil
+
+	// If an error occurs then reopen the storage.
+	if err := func() error {
+		// rowSet maintains the set of rowIDs present in this import.
+		// It allows the cache to be updated once per row, instead of once
+		// per bit.
+		rowSet := make(map[uint64]struct{})
+		lastRowID := uint64(0)
+
+		// Process every bit.
+		for i := range rowIDs {
+			rowID, columnID := rowIDs[i], columnIDs[i]
+
+			// Handle mutex vector (i.e. clear an existing row).
+			if existingRowID, found := f.mutexVector.Get(columnID); found && existingRowID != rowID {
+				// Determine the position of the bit in the storage.
+				pos, err := f.pos(existingRowID, columnID)
+				if err != nil {
+					return err
+				}
+
+				// Clear storage.
+				_, err = f.storage.Remove(pos)
+				if err != nil {
+					return err
+				}
+
+				rowSet[existingRowID] = struct{}{}
+			}
+
+			// Determine the position of the bit in the storage.
+			pos, err := f.pos(rowID, columnID)
+			if err != nil {
+				return err
+			}
+
+			// Write to storage.
+			_, err = f.storage.Add(pos)
+			if err != nil {
+				return err
+			}
+
+			// Reduce the StatsD rate for high volume stats
+			f.stats.Count("ImportBit", 1, 0.0001)
+
+			// Add row to rowSet.
+			if i == 0 || rowID != lastRowID {
+				lastRowID = rowID
+				rowSet[rowID] = struct{}{}
+			}
+
+			// Invalidate block checksum.
+			delete(f.checksums, int(rowID/HashBlockSize))
+		}
+
+		// Update cache counts for all rows.
+		for rowID := range rowSet {
+			// Import should ALWAYS have row() load a new bm from fragment.storage
+			// because the row that's in rowCache hasn't been updated with
+			// this import's data.
+			f.cache.BulkAdd(rowID, f.unprotectedRow(rowID).Count())
+		}
+
+		f.cache.Invalidate()
+
+		return nil
+	}(); err != nil {
+		_ = f.closeStorage()
+		_ = f.openStorage()
+		return err
+	}
+
+	// Write the storage to disk and reload.
+	if err := f.snapshot(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // importValue bulk imports a set of range-encoded values.
 func (f *fragment) importValue(columnIDs, values []uint64, bitDepth uint) error {
 	f.mu.Lock()

fragment_internal_test.go

@@ -1178,6 +1178,76 @@ func TestFragment_SetMutex(t *testing.T) {
 	}
 }
 
+// Ensure a fragment can import mutually exclusive values.
+func TestFragment_ImportMutex(t *testing.T) {
+	tests := []struct {
+		rowIDs []uint64
+		colIDs []uint64
+		exp    map[uint64][]uint64
+	}{
+		{
+			[]uint64{1, 1, 1, 1},
+			[]uint64{0, 1, 2, 3},
+			map[uint64][]uint64{
+				1: {0, 1, 2, 3},
+			},
+		},
+		{
+			[]uint64{1, 1, 1, 1, 2, 2, 2, 2},
+			[]uint64{0, 1, 2, 3, 0, 1, 2, 3},
+			map[uint64][]uint64{
+				1: {},
+				2: {0, 1, 2, 3},
+			},
+		},
+		{
+			[]uint64{1, 1, 1, 1, 2},
+			[]uint64{0, 1, 2, 3, 1},
+			map[uint64][]uint64{
+				1: {0, 2, 3},
+				2: {1},
+			},
+		},
+		{
+			[]uint64{1, 1, 1, 1, 2, 2, 1},
+			[]uint64{0, 1, 2, 3, 1, 8, 1},
+			map[uint64][]uint64{
+				1: {0, 1, 2, 3},
+				2: {8},
+			},
+		},
+		{
+			[]uint64{1, 2, 3},
+			[]uint64{8, 8, 8},
+			map[uint64][]uint64{
+				1: {},
+				2: {},
+				3: {8},
+			},
+		},
+	}
+
+	for i, test := range tests {
+		t.Run(fmt.Sprintf("importmutex%d", i), func(t *testing.T) {
+			f := mustOpenMutexFragment("i", "f", viewStandard, 0, "")
+			defer f.Close()
+
+			err := f.bulkImport(test.rowIDs, test.colIDs)
+			if err != nil {
+				t.Fatalf("bulk importing ids: %v", err)
+			}
+
+			// Check for expected results.
+			for k, v := range test.exp {
+				cols := f.row(k).Columns()
+				if !reflect.DeepEqual(cols, v) {
+					t.Fatalf("expected: %v, but got: %v", v, cols)
+				}
+			}
+		})
+	}
+}
+
 func BenchmarkFragment_Snapshot(b *testing.B) {
 	if *FragmentPath == "" {
 		b.Skip("no fragment specified")