Merge pull request #1532 from travisturner/fragment-rows

Fragment rows() and rowsForColumn()

fragment.go

@@ -47,6 +47,13 @@ const (
 	// ShardWidth is the number of column IDs in a shard.
 	ShardWidth = 1048576
 
+	// containersPerRowSegment is dependent upon ShardWidth,
+	// and it represents the number of containers per shard row
+	// (or rowSegment). Since containers are set in roaring
+	// to be 2^16, then this const should be ShardWidth / 2^16.
+	// It is represented as the exponent n of 2^n.
+	containersPerRowSegment = 4
+
 	// snapshotExt is the file extension used for an in-process snapshot.
 	snapshotExt = ".snapshotting"
 
@@ -153,7 +160,6 @@ func (f *fragment) Open() error {
 		pos := f.storage.Max()
 		f.maxRowID = pos / ShardWidth
 		f.stats.Gauge("rows", float64(f.maxRowID), 1.0)
-
 		return nil
 	}(); err != nil {
 		f.close()
@@ -1680,6 +1686,62 @@ func (f *fragment) readCacheFromArchive(r io.Reader) error {
 	return nil
 }
 
+func (f *fragment) rows() []uint64 {
+	i, _ := f.storage.Containers.Iterator(0)
+	rows := make([]uint64, 0)
+
+	var lastRow uint64
+	lastRow = math.MaxUint64
+
+	// Loop over the existing containers.
+	for i.Next() {
+		key, _ := i.Value()
+
+		// virtual row for the current container
+		vRow := key >> containersPerRowSegment
+
+		// skip dups
+		if vRow == lastRow {
+			continue
+		}
+
+		rows = append(rows, vRow)
+		lastRow = vRow
+	}
+	return rows
+
+}
+
+func (f *fragment) rowsForColumn(columnID uint64) []uint64 {
+	colID := columnID % ShardWidth
+	i, _ := f.storage.Containers.Iterator(0)
+
+	colKey := uint64(0)
+	colVal := uint16(colID & 0xFFFF)
+
+	rows := make([]uint64, 0)
+
+	// Loop over the existing containers.
+	for i.Next() {
+		key, c := i.Value()
+
+		// virtual row for the current container
+		vRow := key >> containersPerRowSegment
+
+		// column container key for virtual row
+		colKey = ((vRow * ShardWidth) + colID) >> 16
+
+		if colKey != key {
+			continue
+		}
+
+		if c.Contains(colVal) {
+			rows = append(rows, vRow)
+		}
+	}
+	return rows
+}
+
 // FragmentBlock represents info about a subsection of the rows in a block.
 // This is used for comparing data in remote blocks for active anti-entropy.
 type FragmentBlock struct {

fragment_internal_test.go

@@ -1278,3 +1278,81 @@ func (f *fragment) mustSetBits(rowID uint64, columnIDs ...uint64) {
 		}
 	}
 }
+
+// Test Various methods of retrieving RowIDs
+func TestFragment_RowsIteration(t *testing.T) {
+	t.Run("firstContainer", func(t *testing.T) {
+		f := mustOpenFragment("i", "f", viewStandard, 0, "")
+		defer f.Close()
+
+		expectedAll := make([]uint64, 0)
+		expectedOdd := make([]uint64, 0)
+		for i := uint64(100); i < uint64(200); i++ {
+			if _, err := f.setBit(i, i%2); err != nil {
+				t.Fatal(err)
+			}
+			expectedAll = append(expectedAll, i)
+			if i%2 == 1 {
+				expectedOdd = append(expectedOdd, i)
+			}
+		}
+
+		ids := f.rows()
+		if !reflect.DeepEqual(expectedAll, ids) {
+			t.Fatalf("Do not match %v %v", expectedAll, ids)
+		}
+
+		ids = f.rowsForColumn(1)
+		if !reflect.DeepEqual(expectedOdd, ids) {
+			t.Fatalf("Do not match %v %v", expectedOdd, ids)
+		}
+	})
+
+	t.Run("secondRow", func(t *testing.T) {
+		f := mustOpenFragment("i", "f", viewStandard, 0, "")
+		defer f.Close()
+
+		expected := []uint64{1, 2}
+		if _, err := f.setBit(1, 66000); err != nil {
+			t.Fatal(err)
+		} else if _, err := f.setBit(2, 66000); err != nil {
+			t.Fatal(err)
+		} else if _, err := f.setBit(2, 166000); err != nil {
+			t.Fatal(err)
+		}
+
+		ids := f.rows()
+		if !reflect.DeepEqual(expected, ids) {
+			t.Fatalf("Do not match %v %v", expected, ids)
+		}
+
+		ids = f.rowsForColumn(66000)
+		if !reflect.DeepEqual(expected, ids) {
+			t.Fatalf("Do not match %v %v", expected, ids)
+		}
+	})
+
+	t.Run("combinations", func(t *testing.T) {
+		f := mustOpenFragment("i", "f", viewStandard, 0, "")
+		defer f.Close()
+
+		expectedRows := make([]uint64, 0)
+		for r := uint64(1); r < uint64(10000); r += 100 {
+			expectedRows = append(expectedRows, r)
+			for c := uint64(1); c < uint64(ShardWidth-1); c += 10000 {
+				if _, err := f.setBit(r, c); err != nil {
+					t.Fatal(err)
+				}
+
+				ids := f.rows()
+				if !reflect.DeepEqual(expectedRows, ids) {
+					t.Fatalf("Do not match %v %v", expectedRows, ids)
+				}
+				ids = f.rowsForColumn(c)
+				if !reflect.DeepEqual(expectedRows, ids) {
+					t.Fatalf("Do not match %v %v", expectedRows, ids)
+				}
+			}
+		}
+	})
+}

roaring/roaring.go

@@ -177,7 +177,7 @@ func (b *Bitmap) Contains(v uint64) bool {
 	if c == nil {
 		return false
 	}
-	return c.contains(lowbits(v))
+	return c.Contains(lowbits(v))
 }
 
 // Remove removes values from the bitmap.
@@ -1272,8 +1272,8 @@ func (c *Container) runAdd(v uint16) bool {
 	return true
 }
 
-// contains returns true if v is in the container.
-func (c *Container) contains(v uint16) bool {
+// Contains returns true if v is in the container.
+func (c *Container) Contains(v uint16) bool {
 	if c.isArray() {
 		return c.arrayContains(v)
 	} else if c.isRun() {