From 619bc1bcd9501157d22c913d71a1d13786a26d25 Mon Sep 17 00:00:00 2001
From: Travis Turner <travis@pilosa.com>
Date: Tue, 25 Sep 2018 14:57:12 -0500
Subject: [PATCH] implements fragment.setRow(row, rowID)

---
 fragment.go               | 54 +++++++++++++++++++++++++++++++++++++--
 fragment_internal_test.go | 48 ++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+), 2 deletions(-)

diff --git a/fragment.go b/fragment.go
index ee586afc6..63c9cce37 100644
--- a/fragment.go
+++ b/fragment.go
@@ -352,11 +352,11 @@ func (f *fragment) unprotectedRow(rowID uint64) *Row {
 	}
 
 	// Only use a subset of the containers.
-	// NOTE: The start & end ranges must be divisible by
+	// NOTE: The start & end ranges must be divisible by container width.
 	data := f.storage.OffsetRange(f.shard*ShardWidth, rowID*ShardWidth, (rowID+1)*ShardWidth)
 
 	// Reference bitmap subrange in storage.
-	// We Clone() data because otherwise row will contains pointers to containers in storage.
+	// We Clone() data because otherwise row will contain pointers to containers in storage.
 	// This causes unexpected results when we cache the row and try to use it later.
 	row := &Row{
 		segments: []rowSegment{{
@@ -491,6 +491,56 @@ func (f *fragment) unprotectedClearBit(rowID, columnID uint64) (changed bool, er
 	return changed, nil
 }
 
+// setRow replaces an existing row (specified by rowID) with the given
+// Row. This updates both the on-disk storage and the in-cache bitmap.
+func (f *fragment) setRow(row *Row, rowID uint64) (bool, error) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.unprotectedSetRow(row, rowID)
+}
+
+func (f *fragment) unprotectedSetRow(row *Row, rowID uint64) (changed bool, err error) {
+	// TODO: In order to return `changed`, we need to first compare
+	// the existing row with the given row. Determine if the overhead
+	// of this is worth having `changed`.
+	// For now we will assume changed is always true.
+	changed = true
+
+	// First container of the row in storage.
+	headContainerKey := rowID << shardVsContainerExponent
+
+	// Remove every existing container in the row.
+	for i := uint64(0); i < (1 << shardVsContainerExponent); i++ {
+		f.storage.Containers.Remove(headContainerKey + i)
+	}
+
+	// From the given row, get the rowSegment for this shard.
+	seg := row.segment(f.shard)
+	if seg == nil {
+		return changed, nil
+	}
+
+	// Put each container from rowSegment to fragment storage.
+	citer, _ := seg.data.Containers.Iterator(f.shard << shardVsContainerExponent)
+	for citer.Next() {
+		k, c := citer.Value()
+		f.storage.Containers.Put(headContainerKey+(k%(1<<shardVsContainerExponent)), c)
+	}
+
+	// Update the row in cache.
+	n := f.storage.CountRange(rowID*ShardWidth, (rowID+1)*ShardWidth)
+	f.cache.BulkAdd(rowID, n)
+
+	// Snapshot storage.
+	if err := f.snapshot(); err != nil {
+		return false, errors.Wrap(err, "snapshotting")
+	}
+
+	f.stats.Count("setRow", 1, 1.0)
+
+	return changed, nil
+}
+
 // ClearRow clears a row for a given rowID within the fragment.
 // This updates both the on-disk storage and the in-cache bitmap.
 func (f *fragment) clearRow(rowID uint64) (bool, error) {
diff --git a/fragment_internal_test.go b/fragment_internal_test.go
index 81ff36859..be0db3b01 100644
--- a/fragment_internal_test.go
+++ b/fragment_internal_test.go
@@ -122,6 +122,54 @@ func TestFragment_ClearRow(t *testing.T) {
 	}
 }
 
+// Ensure a fragment can set a row.
+func TestFragment_SetRow(t *testing.T) {
+	f := mustOpenFragment("i", "f", viewStandard, 7, "")
+	defer f.Close()
+
+	rowID := uint64(1000)
+
+	// Set bits on the fragment.
+	if _, err := f.setBit(rowID, 8000001); err != nil {
+		t.Fatal(err)
+	} else if _, err := f.setBit(rowID, 8065536); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify data on row.
+	if cols := f.row(rowID).Columns(); !reflect.DeepEqual(cols, []uint64{8000001, 8065536}) {
+		t.Fatalf("unexpected columns: %+v", cols)
+	}
+	// Verify count on row.
+	if n := f.row(rowID).Count(); n != 2 {
+		t.Fatalf("unexpected count: %d", n)
+	}
+
+	// Set row (overwrite existing data).
+	row := NewRow(8000002, 8065537, 8131074)
+	if changed, err := f.unprotectedSetRow(row, rowID); err != nil {
+		t.Fatal(err)
+	} else if !changed {
+		t.Fatalf("expected changed value: %v", changed)
+	}
+
+	// Verify data on row.
+	if cols := f.row(rowID).Columns(); !reflect.DeepEqual(cols, []uint64{8000002, 8065537, 8131074}) {
+		t.Fatalf("unexpected columns after set row: %+v", cols)
+	}
+	// Verify count on row.
+	if n := f.row(rowID).Count(); n != 3 {
+		t.Fatalf("unexpected count after set row: %d", n)
+	}
+
+	// Close and reopen the fragment & verify the data.
+	if err := f.reopen(); err != nil {
+		t.Fatal(err)
+	} else if n := f.row(rowID).Count(); n != 3 {
+		t.Fatalf("unexpected count (reopen): %d", n)
+	}
+}
+
 // Ensure a fragment can set & read a value.
 func TestFragment_SetValue(t *testing.T) {
 	t.Run("OK", func(t *testing.T) {