Support 'AND (... OR ... OR ...)' queries

txn.go

@@ -177,6 +177,7 @@ func (txn *Txn) Without(columns ...string) *Txn {
 func (txn *Txn) Union(columns ...string) *Txn {
 	first := !txn.setup
 	txn.initialize()
+
 	for _, columnName := range columns {
 		if idx, ok := txn.columnAt(columnName); ok {
 			txn.rangeReadPair(idx, func(dst, src bitmap.Bitmap) {
@@ -192,6 +193,52 @@ func (txn *Txn) Union(columns ...string) *Txn {
 	return txn
 }
 
+// WithUnion computes a union between all given indexes, and then
+// applies the result to the txn index.
+func (txn *Txn) WithUnion(columns ...string) *Txn {
+	if !txn.setup || len(columns) == 1 {
+		return txn.Union(columns...)
+	}
+
+	// allocate slice of column pointers
+	cols := make([]*column, 0)
+	for _, columnName := range columns {
+		if idx, ok := txn.columnAt(columnName); ok {
+			cols = append(cols, idx)
+		}
+	}
+
+	// allocate temp bitmaps for calculations
+	tmpMap := make(bitmap.Bitmap, 256)
+
+	// adapted from rangeReadPair
+	limit := commit.Chunk(len(txn.index) >> bitmapShift)
+	lock := txn.owner.slock
+
+	// range & lock over each available chunk
+	for chunk := commit.Chunk(0); chunk <= limit; chunk++ {
+		lock.RLock(uint(chunk))
+
+		// reset entire bitmap
+		for i, _ := range tmpMap {
+			tmpMap[i] = 0
+		}
+
+		// for each columm, tmpMap =| colMap
+		for _, orCol := range cols {
+			tmpMap.Or(orCol.Index(chunk))
+		}
+
+		// indexMap =& tmpMap
+		idxMap := chunk.OfBitmap(txn.index)
+		idxMap.And(tmpMap)
+
+		lock.RUnlock(uint(chunk))
+	}
+
+	return txn
+}
+
 // WithValue applies a filter predicate over values for a specific properties. It filters
 // down the items in the query.
 func (txn *Txn) WithValue(column string, predicate func(v interface{}) bool) *Txn {

txn_test.go

@@ -726,6 +726,107 @@ func TestUnion(t *testing.T) {
 	})
 }
 
+func TestWithUnion(t *testing.T) {
+	c := NewCollection()
+	c.CreateColumn("tester", ForString())
+	c.CreateColumn("testerB", ForString())
+
+	c.CreateIndex("tester_1", "tester", func(r Reader) bool { return r.String() == "1" })
+	c.CreateIndex("tester_2", "tester", func(r Reader) bool { return r.String() == "2" })
+	c.CreateIndex("tester_3", "tester", func(r Reader) bool { return r.String() == "3" })
+	c.CreateIndex("testerB_4", "testerB", func(r Reader) bool { return r.String() == "4" })
+	c.CreateIndex("testerB_5", "testerB", func(r Reader) bool { return r.String() == "5" })
+	c.CreateIndex("testerB_6", "testerB", func(r Reader) bool { return r.String() == "6" })
+
+	c.InsertObject(map[string]interface{}{
+		"tester": "1",
+		"testerB": "4",
+	})
+	c.InsertObject(map[string]interface{}{
+		"tester": "2",
+		"testerB": "5",
+	})
+	c.InsertObject(map[string]interface{}{
+		"tester": "3",
+		"testerB": "6",
+	})
+
+	// account for normal use-case
+	c.Query(func(txn *Txn) error {
+		txn.WithUnion("tester_1", "tester_2")
+		txn.Union("testerB_5", "testerB_6")
+
+		assert.Equal(t, 3, txn.Count())
+		return nil
+	})
+
+	// where tester in ['1', '2'] and testerB in ['5', '6']
+	c.Query(func(txn *Txn) error {
+		txn.Union("tester_1", "tester_2")
+		txn.WithUnion("testerB_5", "testerB_6")
+
+		assert.Equal(t, 1, txn.Count())
+		return nil
+	})
+
+	c.Query(func (txn *Txn) error {
+		txn.Without("tester_1", "testerB_5")
+		txn.WithUnion("tester_2", "tester_1", "tester_3")
+
+		assert.Equal(t, 1, txn.Count())
+		return nil
+	})
+}
+
+func TestWithUnionPlayers(t *testing.T) {
+	trueCount := 0
+	players := loadPlayers(100000)
+
+	players.Query(func (txn *Txn) error {
+		ageCol := txn.Any("age")
+		raceCol := txn.Any("race")
+		classCol := txn.Any("class")
+
+		return txn.Range(func (i uint32) {
+			age, _ := ageCol.Get()
+			race, _ := raceCol.Get()
+			class, _ := classCol.Get()
+
+			if race == "dwarf" && (age.(float64) >= 30.0 || class == "mage") {
+				trueCount++
+			}
+		})
+	})
+
+	players.Query(func (txn *Txn) error {
+		txn.With("dwarf")
+		txn.WithUnion("mage", "old")
+
+		assert.Equal(t, trueCount, txn.Count())
+		return nil
+	})
+
+	players.Query(func (txn *Txn) error {
+		txn.With("dwarf", "mage", "old")
+		assert.True(t, txn.Count() < trueCount)
+		return nil
+	})
+
+	players.Query(func (txn *Txn) error {
+		txn.Union("dwarf", "mage", "old")
+		assert.True(t, txn.Count() > trueCount)
+		return nil
+	})
+
+	// dwarf & elf cancel out
+	players.Query(func (txn *Txn) error {
+		txn.With("dwarf")
+		txn.WithUnion("mage", "old", "elf")
+		assert.Equal(t, trueCount, txn.Count())
+		return nil
+	})
+}
+
 func TestSumBalance(t *testing.T) {
 	players := loadPlayers(500)
 	assert.Equal(t, 500, players.Count())