From f0fa03da838bf7e90a491eb95db7dc3975ba6756 Mon Sep 17 00:00:00 2001 From: Ben Johnson Date: Tue, 16 Feb 2016 13:41:37 -0700 Subject: [PATCH] refactor attribute stores This commit refactors the profile and bitmap attribute stores so that they share the same code. There is a new `AttrStore` which associates key/value pairs with a `uint64` identifier. The previous JSON encoding has been fixed to use protobufs which fixes encoding issues for int64. --- attr.go | 210 ++++++++++++++++++++++++++++++++++++++++ attr_test.go | 103 ++++++++++++++++++++ db.go | 128 ++---------------------- db_test.go | 62 ------------ executor.go | 6 +- executor_test.go | 4 +- fragment.go | 8 +- fragment_test.go | 13 +-- frame.go | 170 ++++++-------------------------- frame_test.go | 62 ------------ handler.go | 2 +- handler_test.go | 6 +- internal/internal.pb.go | 18 ++++ internal/internal.proto | 4 + pilosa.go | 50 ---------- 15 files changed, 390 insertions(+), 456 deletions(-) create mode 100644 attr.go create mode 100644 attr_test.go diff --git a/attr.go b/attr.go new file mode 100644 index 000000000..2105bf368 --- /dev/null +++ b/attr.go @@ -0,0 +1,210 @@ +package pilosa + +import ( + "encoding/binary" + "sort" + "sync" + "time" + + "github.com/boltdb/bolt" + "github.com/gogo/protobuf/proto" + "github.com/umbel/pilosa/internal" +) + +// AttrStore represents a storage layer for attributes. +type AttrStore struct { + mu sync.Mutex + path string + db *bolt.DB + + // in-memory cache + attrs map[uint64]map[string]interface{} +} + +// NewAttrStore returns a new instance of AttrStore. +func NewAttrStore(path string) *AttrStore { + return &AttrStore{ + path: path, + attrs: make(map[uint64]map[string]interface{}), + } +} + +// Path returns path to the store's data file. +func (s *AttrStore) Path() string { return s.path } + +// Open opens and initializes the store. +func (s *AttrStore) Open() error { + // Open storage. + db, err := bolt.Open(s.path, 0666, &bolt.Options{Timeout: 1 * time.Second}) + if err != nil { + return err + } + s.db = db + + // Initialize database. + if err := s.db.Update(func(tx *bolt.Tx) error { + if _, err := tx.CreateBucketIfNotExists([]byte("attrs")); err != nil { + return err + } + return nil + }); err != nil { + return err + } + + return nil +} + +// Close closes the store. +func (s *AttrStore) Close() error { return s.db.Close() } + +// Attrs returns a set of attributes by ID. +func (s *AttrStore) Attrs(id uint64) (m map[string]interface{}, err error) { + s.mu.Lock() + defer s.mu.Unlock() + + // Check cache for map. + if m = s.attrs[id]; m != nil { + return m, nil + } + + // Find attributes from storage. + if err = s.db.View(func(tx *bolt.Tx) error { + m, err = txAttrs(tx, id) + if err != nil { + return err + } + return nil + }); err != nil { + return nil, err + } + + // Add to cache. + s.attrs[id] = m + + return +} + +// SetAttrs sets attribute values for a given ID. +func (s *AttrStore) SetAttrs(id uint64, m map[string]interface{}) error { + s.mu.Lock() + defer s.mu.Unlock() + + var attr map[string]interface{} + if err := s.db.Update(func(tx *bolt.Tx) error { + tmp, err := txAttrs(tx, id) + 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) + } else { + attr[k] = 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 + } + + // Swap attributes map in cache. + 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)) + if v == nil { + return emptyMap, nil + } + + var pb internal.AttrMap + if err := proto.Unmarshal(v, &pb); err != nil { + return nil, err + } + return decodeAttrs(pb.GetAttrs()), nil +} + +func encodeAttrs(m map[string]interface{}) []*internal.Attr { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + + a := make([]*internal.Attr, len(keys)) + for i := range keys { + a[i] = encodeAttr(keys[i], m[keys[i]]) + } + return a +} + +func decodeAttrs(pb []*internal.Attr) map[string]interface{} { + m := make(map[string]interface{}, len(pb)) + for i := range pb { + key, value := decodeAttr(pb[i]) + m[key] = value + } + return m +} + +// encodeAttr converts a key/value pair into an Attr internal representation. +func encodeAttr(key string, value interface{}) *internal.Attr { + pb := &internal.Attr{Key: proto.String(key)} + switch value := value.(type) { + case string: + pb.StringValue = proto.String(value) + case float64: + pb.IntValue = proto.Int64(int64(value)) + case int64: + pb.IntValue = proto.Int64(value) + case bool: + pb.BoolValue = proto.Bool(value) + } + return pb +} + +// decodeAttr converts from an Attr internal representation to a key/value pair. +func decodeAttr(attr *internal.Attr) (key string, value interface{}) { + if attr.StringValue != nil { + return attr.GetKey(), attr.GetStringValue() + } else if attr.IntValue != nil { + return attr.GetKey(), attr.GetIntValue() + } else if attr.BoolValue != nil { + return attr.GetKey(), attr.GetBoolValue() + } + return attr.GetKey(), nil +} + +// u64tob encodes v to big endian encoding. +func u64tob(v uint64) []byte { + b := make([]byte, 8) + binary.BigEndian.PutUint64(b, v) + return b +} + +// btou64 decodes b from big endian encoding. +func btou64(b []byte) uint64 { return binary.BigEndian.Uint64(b) } + +// emptyMap is a reusable map that contains no keys. +var emptyMap = make(map[string]interface{}) diff --git a/attr_test.go b/attr_test.go new file mode 100644 index 000000000..af421249c --- /dev/null +++ b/attr_test.go @@ -0,0 +1,103 @@ +package pilosa_test + +import ( + "io/ioutil" + "os" + "reflect" + "testing" + + "github.com/umbel/pilosa" +) + +// Ensure database can set and retrieve profile attributes. +func TestAttrStore_Attrs(t *testing.T) { + s := MustOpenAttrStore() + defer s.Close() + + // Set attributes. + if err := s.SetAttrs(1, map[string]interface{}{"A": int64(100)}); err != nil { + t.Fatal(err) + } else if err := s.SetAttrs(2, map[string]interface{}{"A": int64(200)}); err != nil { + t.Fatal(err) + } else if err := s.SetAttrs(1, map[string]interface{}{"B": "VALUE"}); err != nil { + t.Fatal(err) + } + + // Retrieve attributes for profile #1. + if m, err := s.Attrs(1); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(m, map[string]interface{}{"A": int64(100), "B": "VALUE"}) { + t.Fatalf("unexpected attrs(1): %#v", m) + } + + // Retrieve attributes for profile #2. + if m, err := s.Attrs(2); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(m, map[string]interface{}{"A": int64(200)}) { + t.Fatalf("unexpected attrs(2): %#v", m) + } +} + +// Ensure database returns a non-nil empty map if unset. +func TestAttrStore_Attrs_Empty(t *testing.T) { + s := MustOpenAttrStore() + defer s.Close() + + if m, err := s.Attrs(100); err != nil { + t.Fatal(err) + } else if m == nil || len(m) > 0 { + t.Fatalf("unexpected attrs: %#v", m) + } +} + +// Ensure database can unset attributes if explicitly set to nil. +func TestAttrStore_Attrs_Unset(t *testing.T) { + s := MustOpenAttrStore() + defer s.Close() + + // Set attributes. + if err := s.SetAttrs(1, map[string]interface{}{"A": "X", "B": "Y"}); err != nil { + t.Fatal(err) + } else if err := s.SetAttrs(1, map[string]interface{}{"B": nil}); err != nil { + t.Fatal(err) + } + + // Verify attributes. + if m, err := s.Attrs(1); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(m, map[string]interface{}{"A": "X"}) { + t.Fatalf("unexpected attrs: %#v", m) + } +} + +// AttrStore represents a test wrapper for pilosa.AttrStore. +type AttrStore struct { + *pilosa.AttrStore +} + +// NewAttrStore returns a new instance of AttrStore. +func NewAttrStore() *AttrStore { + f, err := ioutil.TempFile("", "pilosa-attr-") + if err != nil { + panic(err) + } + f.Close() + os.Remove(f.Name()) + + return &AttrStore{AttrStore: pilosa.NewAttrStore(f.Name())} +} + +// MustOpenAttrStore returns a new, opened attribute store at a temporary path. Panic on error. +func MustOpenAttrStore() *AttrStore { + s := NewAttrStore() + if err := s.Open(); err != nil { + panic(err) + } + return s +} + +// Close closes the database and removes the underlying data. +func (s *AttrStore) Close() error { + defer os.RemoveAll(s.Path()) + return s.AttrStore.Close() +} diff --git a/db.go b/db.go index 246d5f289..ca5b2f419 100644 --- a/db.go +++ b/db.go @@ -1,14 +1,10 @@ package pilosa import ( - "encoding/json" "fmt" "os" "path/filepath" "sync" - "time" - - "github.com/boltdb/bolt" ) // DB represents a container for frames. @@ -21,8 +17,7 @@ type DB struct { frames map[string]*Frame // Profile attribute storage and cache - store *bolt.DB - attrs map[uint64]map[string]interface{} + profileAttrStore *AttrStore } // NewDB returns a new instance of DB. @@ -31,7 +26,8 @@ func NewDB(path, name string) *DB { path: path, name: name, frames: make(map[string]*Frame), - attrs: make(map[uint64]map[string]interface{}), + + profileAttrStore: NewAttrStore(filepath.Join(path, "data")), } } @@ -41,6 +37,9 @@ func (db *DB) Name() string { return db.name } // Path returns the path the database was initialized with. func (db *DB) Path() string { return db.path } +// ProfileAttrStore returns the storage for profile attributes. +func (db *DB) ProfileAttrStore() *AttrStore { return db.profileAttrStore } + // Open opens and initializes the database. func (db *DB) Open() error { // Ensure the path exists. @@ -52,7 +51,7 @@ func (db *DB) Open() error { return err } - if err := db.openStore(); err != nil { + if err := db.profileAttrStore.Open(); err != nil { return err } @@ -86,37 +85,14 @@ func (db *DB) openFrames() error { return nil } -// openStore opens and initializes the attribute store. -func (db *DB) openStore() error { - // Open attribute store. - store, err := bolt.Open(filepath.Join(db.path, "data"), 0666, &bolt.Options{Timeout: 1 * time.Second}) - if err != nil { - return err - } - db.store = store - - // Initialize database. - if err := db.store.Update(func(tx *bolt.Tx) error { - if _, err := tx.CreateBucketIfNotExists([]byte("attrs")); err != nil { - return err - } - return nil - }); err != nil { - _ = db.Close() - return err - } - - return nil -} - // Close closes the database and its frames. func (db *DB) Close() error { db.mu.Lock() defer db.mu.Unlock() // Close the attribute store. - if db.store != nil { - db.store.Close() + if db.profileAttrStore != nil { + db.profileAttrStore.Close() } // Close all frames. @@ -176,89 +152,3 @@ func (db *DB) createFrameIfNotExists(name string) (*Frame, error) { return f, nil } - -// ProfileAttrs returns the value of the attribute for a profile. -func (db *DB) ProfileAttrs(id uint64) (m map[string]interface{}, err error) { - db.mu.Lock() - defer db.mu.Unlock() - - // Check cache for map. - if m = db.attrs[id]; m != nil { - return m, nil - } - - // Find attributes from storage. - if err = db.store.View(func(tx *bolt.Tx) error { - m, err = txProfileAttrs(tx, id) - if err != nil { - return err - } - return nil - }); err != nil { - return nil, err - } - - // Add to cache. - db.attrs[id] = m - - return -} - -// SetProfileAttrs sets attribute values for a profile. -func (db *DB) SetProfileAttrs(id uint64, m map[string]interface{}) error { - db.mu.Lock() - defer db.mu.Unlock() - - var attr map[string]interface{} - if err := db.store.Update(func(tx *bolt.Tx) error { - tmp, err := txProfileAttrs(tx, id) - 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) - } else { - attr[k] = v - } - } - - // Marshal and save new values. - buf, err := json.Marshal(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 - } - - // Swap attributes map in cache. - db.attrs[id] = attr - - return nil -} - -// txProfileAttrs returns a map of attributes for a profile. -func txProfileAttrs(tx *bolt.Tx, id uint64) (map[string]interface{}, error) { - if v := tx.Bucket([]byte("attrs")).Get(u64tob(id)); v != nil { - m := make(map[string]interface{}) - if err := json.Unmarshal(v, &m); err != nil { - return nil, err - } - return m, nil - } - return emptyMap, nil -} diff --git a/db_test.go b/db_test.go index d18fb2cb9..5405b07bb 100644 --- a/db_test.go +++ b/db_test.go @@ -3,7 +3,6 @@ package pilosa_test import ( "io/ioutil" "os" - "reflect" "testing" "github.com/umbel/pilosa" @@ -35,67 +34,6 @@ func TestDB_CreateFrameIfNotExists(t *testing.T) { } } -// Ensure database can set and retrieve profile attributes. -func TestDB_ProfileAttrs(t *testing.T) { - db := MustOpenDB() - defer db.Close() - - // Set attributes. - if err := db.SetProfileAttrs(1, map[string]interface{}{"A": float64(100)}); err != nil { - t.Fatal(err) - } else if err := db.SetProfileAttrs(2, map[string]interface{}{"A": float64(200)}); err != nil { - t.Fatal(err) - } else if err := db.SetProfileAttrs(1, map[string]interface{}{"B": "VALUE"}); err != nil { - t.Fatal(err) - } - - // Retrieve attributes for profile #1. - if m, err := db.ProfileAttrs(1); err != nil { - t.Fatal(err) - } else if !reflect.DeepEqual(m, map[string]interface{}{"A": float64(100), "B": "VALUE"}) { - t.Fatalf("unexpected attrs(1): %#v", m) - } - - // Retrieve attributes for profile #2. - if m, err := db.ProfileAttrs(2); err != nil { - t.Fatal(err) - } else if !reflect.DeepEqual(m, map[string]interface{}{"A": float64(200)}) { - t.Fatalf("unexpected attrs(2): %#v", m) - } -} - -// Ensure database returns a non-nil empty map if unset. -func TestDB_ProfileAttrs_Empty(t *testing.T) { - db := MustOpenDB() - defer db.Close() - - if m, err := db.ProfileAttrs(100); err != nil { - t.Fatal(err) - } else if m == nil || len(m) > 0 { - t.Fatalf("unexpected attrs: %#v", m) - } -} - -// Ensure database can unset attributes if explicitly set to nil. -func TestDB_ProfileAttrs_Unset(t *testing.T) { - db := MustOpenDB() - defer db.Close() - - // Set attributes. - if err := db.SetProfileAttrs(1, map[string]interface{}{"A": "X", "B": "Y"}); err != nil { - t.Fatal(err) - } else if err := db.SetProfileAttrs(1, map[string]interface{}{"B": nil}); err != nil { - t.Fatal(err) - } - - // Verify attributes. - if m, err := db.ProfileAttrs(1); err != nil { - t.Fatal(err) - } else if !reflect.DeepEqual(m, map[string]interface{}{"A": "X"}) { - t.Fatalf("unexpected attrs: %#v", m) - } -} - // DB represents a test wrapper for pilosa.DB. type DB struct { *pilosa.DB diff --git a/executor.go b/executor.go index 406af17b3..e94ed677a 100644 --- a/executor.go +++ b/executor.go @@ -117,7 +117,7 @@ func (e *Executor) executeBitmapCall(db string, c pql.BitmapCall, slices []uint6 if c, ok := c.(*pql.Bitmap); ok { fr := e.Index().Frame(db, c.Frame) if fr != nil { - attrs, err := fr.BitmapAttrs(c.ID) + attrs, err := fr.BitmapAttrStore().Attrs(c.ID) if err != nil { return nil, err } @@ -347,7 +347,7 @@ func (e *Executor) executeSetBitmapAttrs(db string, c *pql.SetBitmapAttrs) error } // Set attributes. - if err := frame.SetBitmapAttrs(c.ID, c.Attrs); err != nil { + if err := frame.BitmapAttrStore().SetAttrs(c.ID, c.Attrs); err != nil { return err } @@ -365,7 +365,7 @@ func (e *Executor) executeSetProfileAttrs(db string, c *pql.SetProfileAttrs) err } // Set attributes. - if err := d.SetProfileAttrs(c.ID, c.Attrs); err != nil { + if err := d.ProfileAttrStore().SetAttrs(c.ID, c.Attrs); err != nil { return err } diff --git a/executor_test.go b/executor_test.go index 209c894bd..703e16bc1 100644 --- a/executor_test.go +++ b/executor_test.go @@ -17,7 +17,7 @@ func TestExecutor_Execute_Bitmap(t *testing.T) { idx.MustCreateFragmentIfNotExists("d", "f", 0).MustSetBits(10, 3) idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBits(10, SliceWidth+1) - if err := idx.Frame("d", "f").SetBitmapAttrs(10, map[string]interface{}{"foo": "bar", "baz": 123}); err != nil { + if err := idx.Frame("d", "f").BitmapAttrStore().SetAttrs(10, map[string]interface{}{"foo": "bar", "baz": 123}); err != nil { t.Fatal(err) } @@ -155,7 +155,7 @@ func TestExecutor_Execute_SetBitmapAttrs(t *testing.T) { } f := idx.Frame("d", "f") - if m, err := f.BitmapAttrs(10); err != nil { + if m, err := f.BitmapAttrStore().Attrs(10); err != nil { t.Fatal(err) } else if !reflect.DeepEqual(m, map[string]interface{}{"foo": "bar", "baz": int64(123), "bat": true}) { t.Fatalf("unexpected bitmap attr: %#v", m) diff --git a/fragment.go b/fragment.go index 3173eca52..ed6ec90a3 100644 --- a/fragment.go +++ b/fragment.go @@ -66,10 +66,8 @@ type Fragment struct { LogOutput io.Writer // Bitmap attribute storage. - // Typically this is the parent frame unless overridden for testing. - BitmapAttrStore interface { - BitmapAttrs(id uint64) (map[string]interface{}, error) - } + // This is set by the parent frame unless overridden for testing. + BitmapAttrStore *AttrStore } // NewFragment returns a new instance of Fragment. @@ -403,7 +401,7 @@ func (f *Fragment) TopN(n int, src *Bitmap, field string, fieldValues []interfac // Apply filter, if set. if filters != nil { - attr, err := f.BitmapAttrStore.BitmapAttrs(bitmapID) + attr, err := f.BitmapAttrStore.Attrs(bitmapID) if err != nil { return nil, err } else if attr == nil { diff --git a/fragment_test.go b/fragment_test.go index 3cdb70fce..ce2fc50f7 100644 --- a/fragment_test.go +++ b/fragment_test.go @@ -133,8 +133,8 @@ func TestFragment_TopN_Filter(t *testing.T) { f.MustSetBits(102, 1, 2) // Assign attributes. - f.BitmapAttrStore.SetBitmapAttrs(101, map[string]interface{}{"x": 10}) - f.BitmapAttrStore.SetBitmapAttrs(102, map[string]interface{}{"x": 20}) + f.BitmapAttrStore.SetAttrs(101, map[string]interface{}{"x": 10}) + f.BitmapAttrStore.SetAttrs(102, map[string]interface{}{"x": 20}) // Retrieve top bitmaps. if pairs, err := f.TopN(2, nil, "x", []interface{}{10, 15, 20}); err != nil { @@ -278,7 +278,7 @@ func TestFragment_RankCache_Persistence(t *testing.T) { // Fragment is a test wrapper for pilosa.Fragment. type Fragment struct { *pilosa.Fragment - BitmapAttrStore *BitmapAttrStore + BitmapAttrStore *AttrStore } // NewFragment returns a new instance of Fragment with a temporary path. @@ -291,9 +291,9 @@ func NewFragment(db, frame string, slice uint64) *Fragment { f := &Fragment{ Fragment: pilosa.NewFragment(file.Name(), db, frame, slice), - BitmapAttrStore: NewBitmapAttrStore(), + BitmapAttrStore: MustOpenAttrStore(), } - f.Fragment.BitmapAttrStore = f.BitmapAttrStore + f.Fragment.BitmapAttrStore = f.BitmapAttrStore.AttrStore return f } @@ -310,6 +310,7 @@ func MustOpenFragment(db, frame string, slice uint64) *Fragment { func (f *Fragment) Close() error { defer os.Remove(f.Path()) defer os.Remove(f.CachePath()) + defer f.BitmapAttrStore.Close() return f.Fragment.Close() } @@ -321,7 +322,7 @@ func (f *Fragment) Reopen() error { } f.Fragment = pilosa.NewFragment(path, f.DB(), f.Frame(), f.Slice()) - f.Fragment.BitmapAttrStore = f.BitmapAttrStore + f.Fragment.BitmapAttrStore = f.BitmapAttrStore.AttrStore if err := f.Open(); err != nil { return err } diff --git a/frame.go b/frame.go index f37ce0b3e..6142f5ec7 100644 --- a/frame.go +++ b/frame.go @@ -1,16 +1,11 @@ package pilosa import ( - "encoding/binary" - "encoding/json" "fmt" "os" "path/filepath" "strconv" "sync" - "time" - - "github.com/boltdb/bolt" ) // Frame represents a container for fragments. @@ -24,8 +19,7 @@ type Frame struct { fragments map[uint64]*Fragment // Bitmap attribute storage and cache - store *bolt.DB - attrs map[uint64]map[string]interface{} + bitmapAttrStore *AttrStore } // NewFrame returns a new instance of frame. @@ -35,8 +29,8 @@ func NewFrame(path, db, name string) *Frame { db: db, name: name, - fragments: make(map[uint64]*Fragment), - attrs: make(map[uint64]map[string]interface{}), + fragments: make(map[uint64]*Fragment), + bitmapAttrStore: NewAttrStore(filepath.Join(path, "data")), } } @@ -49,6 +43,9 @@ func (f *Frame) DB() string { return f.db } // Path returns the path the frame was initialized with. func (f *Frame) Path() string { return f.path } +// BitmapAttrStore returns the attribute storage. +func (f *Frame) BitmapAttrStore() *AttrStore { return f.bitmapAttrStore } + // SliceN returns the max slice in the frame. func (f *Frame) SliceN() uint64 { f.mu.Lock() @@ -65,16 +62,23 @@ func (f *Frame) SliceN() uint64 { // Open opens and initializes the frame. func (f *Frame) Open() error { - // Ensure the frame's path exists. - if err := os.MkdirAll(f.path, 0777); err != nil { - return err - } + if err := func() error { + // Ensure the frame's path exists. + if err := os.MkdirAll(f.path, 0777); err != nil { + return err + } - if err := f.openFragments(); err != nil { - return err - } + if err := f.openFragments(); err != nil { + return err + } - if err := f.openStore(); err != nil { + if err := f.bitmapAttrStore.Open(); err != nil { + return err + } + + return nil + }(); err != nil { + f.Close() return err } @@ -109,44 +113,21 @@ func (f *Frame) openFragments() error { if err := frag.Open(); err != nil { return fmt.Errorf("open fragment: slice=%s, err=%s", frag.Slice(), err) } - frag.BitmapAttrStore = f + frag.BitmapAttrStore = f.bitmapAttrStore f.fragments[frag.Slice()] = frag } return nil } -// openStore opens and initializes the attribute store. -func (f *Frame) openStore() error { - // Open attribute store. - store, err := bolt.Open(filepath.Join(f.path, "data"), 0666, &bolt.Options{Timeout: 1 * time.Second}) - if err != nil { - return err - } - f.store = store - - // Initialize database. - if err := f.store.Update(func(tx *bolt.Tx) error { - if _, err := tx.CreateBucketIfNotExists([]byte("attrs")); err != nil { - return err - } - return nil - }); err != nil { - _ = f.Close() - return err - } - - return nil -} - // Close closes the frame and its fragments. func (f *Frame) Close() error { f.mu.Lock() defer f.mu.Unlock() // Close the attribute store. - if f.store != nil { - _ = f.store.Close() + if f.bitmapAttrStore != nil { + _ = f.bitmapAttrStore.Close() } // Close all fragments. @@ -190,107 +171,10 @@ func (f *Frame) createFragmentIfNotExists(slice uint64) (*Fragment, error) { if err := frag.Open(); err != nil { return nil, err } - frag.BitmapAttrStore = f + frag.BitmapAttrStore = f.bitmapAttrStore + + // Save to lookup. f.fragments[slice] = frag return frag, nil } - -// BitmapAttrs returns the value of the attribute for a bitmap. -func (f *Frame) BitmapAttrs(id uint64) (m map[string]interface{}, err error) { - f.mu.Lock() - defer f.mu.Unlock() - - // Check cache for map. - if m = f.attrs[id]; m != nil { - return m, nil - } - - // Find attributes from storage. - if err = f.store.View(func(tx *bolt.Tx) error { - m, err = txBitmapAttrs(tx, id) - if err != nil { - return err - } - return nil - }); err != nil { - return nil, err - } - - // Add to cache. - f.attrs[id] = m - - return -} - -// SetBitmapAttrs sets attribute values for a bitmap. -func (f *Frame) SetBitmapAttrs(id uint64, m map[string]interface{}) error { - f.mu.Lock() - defer f.mu.Unlock() - - var attr map[string]interface{} - if err := f.store.Update(func(tx *bolt.Tx) error { - tmp, err := txBitmapAttrs(tx, id) - 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) - } else { - attr[k] = v - } - } - - // Marshal and save new values. - buf, err := json.Marshal(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 - } - - // Swap attributes map in cache. - f.attrs[id] = attr - - return nil -} - -// txBitmapAttrs returns a map of attributes for a bitmap. -func txBitmapAttrs(tx *bolt.Tx, id uint64) (map[string]interface{}, error) { - if v := tx.Bucket([]byte("attrs")).Get(u64tob(id)); v != nil { - m := make(map[string]interface{}) - if err := json.Unmarshal(v, &m); err != nil { - return nil, err - } - return m, nil - } - return emptyMap, nil -} - -// u64tob encodes v to big endian encoding. -func u64tob(v uint64) []byte { - b := make([]byte, 8) - binary.BigEndian.PutUint64(b, v) - return b -} - -// btou64 decodes b from big endian encoding. -func btou64(b []byte) uint64 { return binary.BigEndian.Uint64(b) } - -// emptyMap is a reusable map that contains no keys. -var emptyMap = make(map[string]interface{}) diff --git a/frame_test.go b/frame_test.go index 391004982..60eab3113 100644 --- a/frame_test.go +++ b/frame_test.go @@ -3,7 +3,6 @@ package pilosa_test import ( "io/ioutil" "os" - "reflect" "testing" "github.com/umbel/pilosa" @@ -35,67 +34,6 @@ func TestFrame_CreateFragmentIfNotExists(t *testing.T) { } } -// Ensure frame can set and retrieve bitmap attributes. -func TestFrame_BitmapAttrs(t *testing.T) { - f := MustOpenFrame() - defer f.Close() - - // Set attributes. - if err := f.SetBitmapAttrs(1, map[string]interface{}{"A": float64(100)}); err != nil { - t.Fatal(err) - } else if err := f.SetBitmapAttrs(2, map[string]interface{}{"A": float64(200)}); err != nil { - t.Fatal(err) - } else if err := f.SetBitmapAttrs(1, map[string]interface{}{"B": "VALUE"}); err != nil { - t.Fatal(err) - } - - // Retrieve attributes for bitmap #1. - if m, err := f.BitmapAttrs(1); err != nil { - t.Fatal(err) - } else if !reflect.DeepEqual(m, map[string]interface{}{"A": float64(100), "B": "VALUE"}) { - t.Fatalf("unexpected attrs(1): %#v", m) - } - - // Retrieve attributes for bitmap #2. - if m, err := f.BitmapAttrs(2); err != nil { - t.Fatal(err) - } else if !reflect.DeepEqual(m, map[string]interface{}{"A": float64(200)}) { - t.Fatalf("unexpected attrs(2): %#v", m) - } -} - -// Ensure frame returns a non-nil empty map if unset. -func TestFrame_BitmapAttrs_Empty(t *testing.T) { - f := MustOpenFrame() - defer f.Close() - - if m, err := f.BitmapAttrs(100); err != nil { - t.Fatal(err) - } else if m == nil || len(m) > 0 { - t.Fatalf("unexpected attrs: %#v", m) - } -} - -// Ensure frame can unset attributes if explicitly set to nil. -func TestFrame_BitmapAttrs_Unset(t *testing.T) { - f := MustOpenFrame() - defer f.Close() - - // Set attributes. - if err := f.SetBitmapAttrs(1, map[string]interface{}{"A": "X", "B": "Y"}); err != nil { - t.Fatal(err) - } else if err := f.SetBitmapAttrs(1, map[string]interface{}{"B": nil}); err != nil { - t.Fatal(err) - } - - // Verify attributes. - if m, err := f.BitmapAttrs(1); err != nil { - t.Fatal(err) - } else if !reflect.DeepEqual(m, map[string]interface{}{"A": "X"}) { - t.Fatalf("unexpected attrs: %#v", m) - } -} - // Frame represents a test wrapper for pilosa.Frame. type Frame struct { *pilosa.Frame diff --git a/handler.go b/handler.go index 5a05200b9..c6927a14e 100644 --- a/handler.go +++ b/handler.go @@ -138,7 +138,7 @@ func (h *Handler) readProfiles(db *DB, ids []uint64) ([]*Profile, error) { a := make([]*Profile, 0, len(ids)) for _, id := range ids { // Read attributes for profile. Skip profile if empty. - attrs, err := db.ProfileAttrs(id) + attrs, err := db.ProfileAttrStore().Attrs(id) if err != nil { return nil, err } else if len(attrs) == 0 { diff --git a/handler_test.go b/handler_test.go index 04fec304b..ebb9cf78b 100644 --- a/handler_test.go +++ b/handler_test.go @@ -161,9 +161,9 @@ func TestHandler_Query_Bitmap_Profiles_JSON(t *testing.T) { db, err := idx.CreateDBIfNotExists("d") if err != nil { t.Fatal(err) - } else if err := db.SetProfileAttrs(3, map[string]interface{}{"x": "y"}); err != nil { + } else if err := db.ProfileAttrStore().SetAttrs(3, map[string]interface{}{"x": "y"}); err != nil { t.Fatal(err) - } else if err := db.SetProfileAttrs(66, map[string]interface{}{"y": 123, "z": false}); err != nil { + } else if err := db.ProfileAttrStore().SetAttrs(66, map[string]interface{}{"y": 123, "z": false}); err != nil { t.Fatal(err) } @@ -226,7 +226,7 @@ func TestHandler_Query_Bitmap_Profiles_Protobuf(t *testing.T) { db, err := idx.CreateDBIfNotExists("d") if err != nil { t.Fatal(err) - } else if err := db.SetProfileAttrs(1, map[string]interface{}{"x": "y"}); err != nil { + } else if err := db.ProfileAttrStore().SetAttrs(1, map[string]interface{}{"x": "y"}); err != nil { t.Fatal(err) } diff --git a/internal/internal.pb.go b/internal/internal.pb.go index 67a882094..cb128d7c6 100644 --- a/internal/internal.pb.go +++ b/internal/internal.pb.go @@ -15,6 +15,7 @@ It has these top-level messages: Bit Profile Attr + AttrMap QueryRequest QueryResponse ImportRequest @@ -192,6 +193,22 @@ func (m *Attr) GetBoolValue() bool { return false } +type AttrMap struct { + Attrs []*Attr `protobuf:"bytes,1,rep,name=Attrs" json:"Attrs,omitempty"` + XXX_unrecognized []byte `json:"-"` +} + +func (m *AttrMap) Reset() { *m = AttrMap{} } +func (m *AttrMap) String() string { return proto.CompactTextString(m) } +func (*AttrMap) ProtoMessage() {} + +func (m *AttrMap) GetAttrs() []*Attr { + if m != nil { + return m.Attrs + } + return nil +} + type QueryRequest struct { DB *string `protobuf:"bytes,1,req,name=DB" json:"DB,omitempty"` Query *string `protobuf:"bytes,2,req,name=Query" json:"Query,omitempty"` @@ -367,6 +384,7 @@ func init() { proto.RegisterType((*Bit)(nil), "internal.Bit") proto.RegisterType((*Profile)(nil), "internal.Profile") proto.RegisterType((*Attr)(nil), "internal.Attr") + proto.RegisterType((*AttrMap)(nil), "internal.AttrMap") proto.RegisterType((*QueryRequest)(nil), "internal.QueryRequest") proto.RegisterType((*QueryResponse)(nil), "internal.QueryResponse") proto.RegisterType((*ImportRequest)(nil), "internal.ImportRequest") diff --git a/internal/internal.proto b/internal/internal.proto index 9de514fb6..f40e8d61b 100644 --- a/internal/internal.proto +++ b/internal/internal.proto @@ -32,6 +32,10 @@ message Attr { optional bool BoolValue = 4; } +message AttrMap { + repeated Attr Attrs = 1; +} + message QueryRequest { required string DB = 1; required string Query = 2; diff --git a/pilosa.go b/pilosa.go index a6767dc11..f788a2037 100644 --- a/pilosa.go +++ b/pilosa.go @@ -2,7 +2,6 @@ package pilosa import ( "errors" - "sort" "github.com/gogo/protobuf/proto" "github.com/umbel/pilosa/internal" @@ -73,52 +72,3 @@ func decodeProfile(pb *internal.Profile) *Profile { return p } - -func encodeAttrs(m map[string]interface{}) []*internal.Attr { - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - sort.Strings(keys) - - a := make([]*internal.Attr, len(keys)) - for i := range keys { - a[i] = encodeAttr(keys[i], m[keys[i]]) - } - return a -} - -func decodeAttrs(pb []*internal.Attr) map[string]interface{} { - m := make(map[string]interface{}, len(pb)) - for i := range pb { - key, value := decodeAttr(pb[i]) - m[key] = value - } - return m -} - -// encodeAttr converts a key/value pair into an Attr internal representation. -func encodeAttr(key string, value interface{}) *internal.Attr { - pb := &internal.Attr{Key: proto.String(key)} - switch value := value.(type) { - case string: - pb.StringValue = proto.String(value) - case int64: - pb.IntValue = proto.Int64(value) - case bool: - pb.BoolValue = proto.Bool(value) - } - return pb -} - -// decodeAttr converts from an Attr internal representation to a key/value pair. -func decodeAttr(attr *internal.Attr) (key string, value interface{}) { - if attr.StringValue != nil { - return attr.GetKey(), attr.GetStringValue() - } else if attr.IntValue != nil { - return attr.GetKey(), attr.GetIntValue() - } else if attr.BoolValue != nil { - return attr.GetKey(), attr.GetBoolValue() - } - return attr.GetKey(), nil -}