Merge pull request #93 from benbjohnson/set-bulk-attrs

Optimize bulk bitmap attribute update

attr.go

@@ -97,47 +97,12 @@ func (s *AttrStore) SetAttrs(id uint64, m map[string]interface{}) error {
 
 	var attr map[string]interface{}
 	if err := s.db.Update(func(tx *bolt.Tx) error {
-		tmp, err := txAttrs(tx, id)
+		tmp, err := txUpdateAttrs(tx, id, m)
 		if err != nil {
 			return err
 		}
 		attr = tmp
 
-		// Create a new map if it is empty so we don't update emptyMap.
-		if len(attr) == 0 {
-			attr = make(map[string]interface{}, len(m))
-		}
-
-		// Merge attributes with original values.
-		// Nil values should delete keys.
-		for k, v := range m {
-			if v == nil {
-				delete(attr, k)
-				continue
-			}
-
-			switch v := v.(type) {
-			case int:
-				attr[k] = uint64(v)
-			case uint:
-				attr[k] = uint64(v)
-			case int64:
-				attr[k] = uint64(v)
-			case string, uint64, bool:
-				attr[k] = v
-			default:
-				return fmt.Errorf("invalid attr type: %T", v)
-			}
-		}
-
-		// Marshal and save new values.
-		buf, err := proto.Marshal(&internal.AttrMap{Attrs: encodeAttrs(attr)})
-		if err != nil {
-			return err
-		}
-		if err := tx.Bucket([]byte("attrs")).Put(u64tob(id), buf); err != nil {
-			return err
-		}
 		return nil
 	}); err != nil {
 		return err
@@ -149,6 +114,42 @@ func (s *AttrStore) SetAttrs(id uint64, m map[string]interface{}) error {
 	return nil
 }
 
+// SetBulkAttrs sets attribute values for a set of ids.
+func (s *AttrStore) SetBulkAttrs(m map[uint64]map[string]interface{}) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	attrs := make(map[uint64]map[string]interface{})
+	if err := s.db.Update(func(tx *bolt.Tx) error {
+		// Collect and sort keys.
+		ids := make([]uint64, 0, len(m))
+		for id := range m {
+			ids = append(ids, id)
+		}
+		sort.Sort(uint64Slice(ids))
+
+		// Update attributes for each id.
+		for _, id := range ids {
+			attr, err := txUpdateAttrs(tx, id, m[id])
+			if err != nil {
+				return err
+			}
+			attrs[id] = attr
+		}
+
+		return nil
+	}); err != nil {
+		return err
+	}
+
+	// Swap attributes map in cache.
+	for id, attr := range attrs {
+		s.attrs[id] = attr
+	}
+
+	return nil
+}
+
 // txAttrs returns a map of attributes for a bitmap.
 func txAttrs(tx *bolt.Tx, id uint64) (map[string]interface{}, error) {
 	v := tx.Bucket([]byte("attrs")).Get(u64tob(id))
@@ -163,6 +164,52 @@ func txAttrs(tx *bolt.Tx, id uint64) (map[string]interface{}, error) {
 	return decodeAttrs(pb.GetAttrs()), nil
 }
 
+// txUpdateAttrs updates the attributes for an id.
+// Returns the new combined set of attributes for the id.
+func txUpdateAttrs(tx *bolt.Tx, id uint64, m map[string]interface{}) (map[string]interface{}, error) {
+	attr, err := txAttrs(tx, id)
+	if err != nil {
+		return nil, err
+	}
+
+	// Create a new map if it is empty so we don't update emptyMap.
+	if len(attr) == 0 {
+		attr = make(map[string]interface{}, len(m))
+	}
+
+	// Merge attributes with original values.
+	// Nil values should delete keys.
+	for k, v := range m {
+		if v == nil {
+			delete(attr, k)
+			continue
+		}
+
+		switch v := v.(type) {
+		case int:
+			attr[k] = uint64(v)
+		case uint:
+			attr[k] = uint64(v)
+		case int64:
+			attr[k] = uint64(v)
+		case string, uint64, bool:
+			attr[k] = v
+		default:
+			return nil, fmt.Errorf("invalid attr type: %T", v)
+		}
+	}
+
+	// Marshal and save new values.
+	buf, err := proto.Marshal(&internal.AttrMap{Attrs: encodeAttrs(attr)})
+	if err != nil {
+		return nil, err
+	}
+	if err := tx.Bucket([]byte("attrs")).Put(u64tob(id), buf); err != nil {
+		return nil, err
+	}
+	return attr, nil
+}
+
 func encodeAttrs(m map[string]interface{}) []*internal.Attr {
 	keys := make([]string, 0, len(m))
 	for k := range m {
@@ -216,6 +263,15 @@ func decodeAttr(attr *internal.Attr) (key string, value interface{}) {
 	return attr.GetKey(), nil
 }
 
+// cloneAttrs returns a shallow clone of m.
+func cloneAttrs(m map[string]interface{}) map[string]interface{} {
+	other := make(map[string]interface{}, len(m))
+	for k, v := range m {
+		other[k] = v
+	}
+	return other
+}
+
 // u64tob encodes v to big endian encoding.
 func u64tob(v uint64) []byte {
 	b := make([]byte, 8)

executor.go

@@ -62,6 +62,11 @@ func (e *Executor) Execute(db string, q *pql.Query, slices []uint64, opt *ExecOp
 		}
 	}
 
+	// Optimize handling for bulk attribute insertion.
+	if hasOnlySetBitmapAttrs(q.Calls) {
+		return e.executeBulkSetBitmapAttrs(db, q.Calls)
+	}
+
 	// Execute each call serially.
 	results := make([]interface{}, 0, len(q.Calls))
 	for _, call := range q.Calls {
@@ -448,6 +453,51 @@ func (e *Executor) executeSetBitmapAttrs(db string, c *pql.SetBitmapAttrs) error
 	return nil
 }
 
+// executeBulkSetBitmapAttrs executes a set of SetBitmapAttrs() calls.
+func (e *Executor) executeBulkSetBitmapAttrs(db string, calls pql.Calls) ([]interface{}, error) {
+	// Collect attributes by frame/id.
+	m := make(map[string]map[uint64]map[string]interface{})
+	for _, call := range calls {
+		c := call.(*pql.SetBitmapAttrs)
+
+		// Create frame group, if not exists.
+		frameMap := m[c.Frame]
+		if frameMap == nil {
+			frameMap = make(map[uint64]map[string]interface{})
+			m[c.Frame] = frameMap
+		}
+
+		// Set or merge attributes.
+		attr := frameMap[c.ID]
+		if attr == nil {
+			frameMap[c.ID] = cloneAttrs(c.Attrs)
+		} else {
+			for k, v := range c.Attrs {
+				attr[k] = v
+			}
+		}
+	}
+
+	// Bulk insert attributes by frame.
+	for name, frameMap := range m {
+		// Retrieve frame.
+		frame, err := e.Index.CreateFrameIfNotExists(db, name)
+		if err != nil {
+			return nil, err
+		}
+
+		// Set attributes.
+		if err := frame.BitmapAttrStore().SetBulkAttrs(frameMap); err != nil {
+			return nil, err
+		}
+	}
+
+	// TODO: Propagate attributes to other servers in cluster.
+
+	// Return a set of nil responses to match the non-optimized return.
+	return make([]interface{}, len(calls)), nil
+}
+
 // executeSetProfileAttrs executes a SetProfileAttrs() call.
 func (e *Executor) executeSetProfileAttrs(db string, c *pql.SetProfileAttrs) error {
 	// Retrieve database.
@@ -583,3 +633,17 @@ func decodeError(s string) error {
 	}
 	return errors.New(s)
 }
+
+// hasOnlySetBitmapAttrs returns true if calls only contains SetBitmapAttrs() calls.
+func hasOnlySetBitmapAttrs(calls pql.Calls) bool {
+	if len(calls) == 0 {
+		return false
+	}
+
+	for _, call := range calls {
+		if _, ok := call.(*pql.SetBitmapAttrs); !ok {
+			return false
+		}
+	}
+	return true
+}