go/zodb/btree: New package to work with ZODB BTrees (draft)

Provide minimal support for BTrees.LOBTree Get for now.

go/zodb/btree/btree.go

+// Copyright (c) 2001, 2002 Zope Foundation and Contributors.
+// All Rights Reserved.
+//
+// Copyright (C) 2018  Nexedi SA and Contributors.
+//                     Kirill Smelkov <kirr@nexedi.com>
+//
+// This software is subject to the provisions of the Zope Public License,
+// Version 2.1 (ZPL).  A copy of the ZPL should accompany this distribution.
+// THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED
+// WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+// WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS
+// FOR A PARTICULAR PURPOSE
+
+// Package btree provides B⁺ Trees for ZODB.
+//
+// It is modelled and data compatible with BTree/py package:
+//
+//	http://btrees.readthedocs.io
+//	https://github.com/zopefoundation/BTrees
+package btree
+
+// See https://github.com/zopefoundation/BTrees/blob/4.5.0-1-gc8bf24e/BTrees/Development.txt#L198
+// for BTree & Bucket organization details.
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+	"sort"
+
+	pickle "github.com/kisielk/og-rek"
+	"lab.nexedi.com/kirr/go123/xerr"
+	"lab.nexedi.com/kirr/neo/go/zodb"
+)
+
+// KEY is the type for BTree keys.
+//
+// XXX -> template?
+type KEY int64
+
+// Bucket is a leaf node of a B⁺ tree.
+//
+// It mimics ?OBucket from btree/py, with ? being any integer.
+type Bucket struct {
+	zodb.Persistent
+
+	// https://github.com/zopefoundation/BTrees/blob/4.5.0-1-gc8bf24e/BTrees/BTreeModuleTemplate.c#L179:
+	//
+	// A Bucket wraps contiguous vectors of keys and values.  Keys are unique,
+	// and stored in sorted order.  The 'values' pointer may be NULL if the
+	// Bucket is used to implement a set.  Buckets serving as leafs of BTrees
+	// are chained together via 'next', so that the entire BTree contents
+	// can be traversed in sorted order quickly and easily.
+
+	next   *Bucket       // the bucket with the next-larger keys
+	keys   []KEY         // 'len' keys, in increasing order
+	values []interface{} // 'len' corresponding values
+}
+
+// _BTreeItem mimics BTreeItem from btree/py.
+//
+// XXX export for BTree.Children?
+type _BTreeItem struct {
+	key   KEY
+	child interface{} // BTree or Bucket
+}
+
+// BTree is a non-leaf node of a B⁺ tree.
+//
+// It mimics ?OBTree from btree/py, with ? being any integer.
+type BTree struct {
+	zodb.Persistent
+
+	// https://github.com/zopefoundation/BTrees/blob/4.5.0-1-gc8bf24e/BTrees/BTreeModuleTemplate.c#L205:
+	//
+	// firstbucket points to the bucket containing the smallest key in
+	// the BTree.  This is found by traversing leftmost child pointers
+	// (data[0].child) until reaching a Bucket.
+	firstbucket *Bucket
+
+	// https://github.com/zopefoundation/BTrees/blob/4.5.0-1-gc8bf24e/BTrees/BTreeModuleTemplate.c#L211:
+	//
+	// The BTree points to 'len' children, via the "child" fields of the data
+	// array.  There are len-1 keys in the 'key' fields, stored in increasing
+	// order.  data[0].key is unused.  For i in 0 .. len-1, all keys reachable
+	// from data[i].child are >= data[i].key and < data[i+1].key, at the
+	// endpoints pretending that data[0].key is -∞ and data[len].key is +∞.
+	data []_BTreeItem
+}
+
+// Get searches BTree by key.
+//
+// It loads intermediate BTree nodes from database on demand as needed.
+//
+// t need not be activated beforehand for Get to work.
+func (t *BTree) Get(ctx context.Context, key KEY) (_ interface{}, _ bool, err error) {
+	defer xerr.Contextf(&err, "btree(%s): get %v", t.POid(), key)
+	err = t.PActivate(ctx)
+	if err != nil {
+		return nil, false, err
+	}
+
+	if len(t.data) == 0 {
+		// empty btree
+		t.PDeactivate()
+		return nil, false, nil
+	}
+
+	for {
+		// search i: K(i) ≤ k < K(i+1)	; K(0) = -∞, K(len) = +∞
+		i := sort.Search(len(t.data), func(i int) bool {
+			j := i + 1
+			if j == len(t.data) {
+				return true // [len].key = +∞
+			}
+			return key < t.data[j].key
+		})
+
+		switch child := t.data[i].child.(type) {
+		case *BTree:
+			t.PDeactivate()
+			t = child
+			err = t.PActivate(ctx)
+			if err != nil {
+				return nil, false, err
+			}
+
+		case *Bucket:
+			t.PDeactivate()
+			return child.Get(ctx, key)
+		}
+	}
+}
+
+// Get searches Bucket by key.
+//
+// TODO Bucket.Get should not get ctx argument and just require that the bucket
+// must be already activated. Correspondingly there should be no error returned.
+func (b *Bucket) Get(ctx context.Context, key KEY) (_ interface{}, _ bool, err error) {
+	defer xerr.Contextf(&err, "bucket(%s): get %v", b.POid(), key)
+	err = b.PActivate(ctx)
+	if err != nil {
+		return nil, false, err
+	}
+
+	v, ok := b.get(key)
+	b.PDeactivate()
+	return v, ok, nil
+}
+
+// get searches Bucket by key.
+//
+// no loading from database is done. The bucket must be already activated.
+func (b *Bucket) get(key KEY) (interface{}, bool) {
+	// search i: K(i-1) < k ≤ K(i)		; K(-1) = -∞, K(len) = +∞
+	i := sort.Search(len(b.keys), func(i int) bool {
+		return key <= b.keys[i]
+	})
+
+	if i == len(b.keys) || b.keys[i] != key {
+		return nil, false // not found
+	}
+	return b.values[i], true
+}
+
+// TODO Bucket.MinKey
+// TODO Bucket.MaxKey
+
+
+// ---- serialization ----
+
+// https://github.com/zopefoundation/BTrees/blob/4.5.0-1-gc8bf24e/BTrees/BucketTemplate.c#L1195:
+//
+// For a mapping bucket (self->values is not NULL), a one-tuple or two-tuple.
+// The first element is a tuple interleaving keys and values, of length
+// 2 * self->len.  The second element is the next bucket, present iff next is
+// non-NULL:
+//
+//	(
+//	     (keys[0], values[0], keys[1], values[1], ...,
+//	                          keys[len-1], values[len-1]),
+//	     <self->next iff non-NULL>
+//	)
+
+type bucketState Bucket // hide state methods from public API
+
+// DropState implements Stateful to discard bucket state.
+func (b *bucketState) DropState() {
+	b.next = nil
+	b.keys = nil
+	b.values = nil
+}
+
+// PySetState implements PyStateful to set bucket data from pystate.
+func (b *bucketState) PySetState(pystate interface{}) (err error) {
+	t, ok := pystate.(pickle.Tuple)
+	if !ok {
+		return fmt.Errorf("top: expect (...); got %T", pystate)
+	}
+	if !(1 <= len(t) && len(t) <= 2) {
+		return fmt.Errorf("top: expect [1..2](); got [%d]()", len(t))
+	}
+
+	// .next present
+	if len(t) == 2 {
+		next, ok := t[1].(*Bucket)
+		if !ok {
+			return fmt.Errorf(".next must be Bucket; got %T", t[1])
+		}
+		b.next = next
+	}
+
+	// main part
+	t, ok = t[0].(pickle.Tuple)
+	if !ok {
+		return fmt.Errorf("data: expect (...); got %T", t[0])
+	}
+	if len(t)%2 != 0 {
+		return fmt.Errorf("data: expect [%%2](); got [%d]()", len(t))
+	}
+
+	// reset arrays just in case
+	n := len(t) / 2
+	b.keys = make([]KEY, 0, n)
+	b.values = make([]interface{}, 0, n)
+
+	for i := 0; i < n; i++ {
+		xk := t[2*i]
+		v := t[2*i+1]
+
+		k, ok := xk.(int64) // XXX use KEY	XXX -> Xint64
+		if !ok {
+			return fmt.Errorf("data: [%d]: key must be integer; got %T", i, xk)
+		}
+
+		// XXX check keys are sorted?
+		b.keys = append(b.keys, KEY(k)) // XXX cast
+		b.values = append(b.values, v)
+	}
+
+	return nil
+}
+
+
+// https://github.com/zopefoundation/BTrees/blob/4.5.0-1-gc8bf24e/BTrees/BTreeTemplate.c#L1087:
+//
+// For an empty BTree (self->len == 0), None.
+//
+// For a BTree with one child (self->len == 1), and that child is a bucket,
+// and that bucket has a NULL oid, a one-tuple containing a one-tuple
+// containing the bucket's state:
+//
+//     (
+//         (
+//              child[0].__getstate__(),
+//         ),
+//     )
+//
+// Else a two-tuple.  The first element is a tuple interleaving the BTree's
+// keys and direct children, of size 2*self->len - 1 (key[0] is unused and
+// is not saved).  The second element is the firstbucket:
+//
+//     (
+//          (child[0], key[1], child[1], key[2], child[2], ...,
+//                                       key[len-1], child[len-1]),
+//          self->firstbucket
+//     )
+//
+// In the above, key[i] means self->data[i].key, and similarly for child[i].
+
+type btreeState BTree // hide state methods from public API
+
+// DropState implements zodb.Stateful.
+func (t *btreeState) DropState() {
+	t.firstbucket = nil
+	t.data = nil
+}
+
+// PySetState implements zodb.PyStateful to set btree data from pystate.
+func (bt *btreeState) PySetState(pystate interface{}) (err error) {
+	// empty btree
+	if _, ok := pystate.(pickle.None); ok {
+		bt.firstbucket = nil
+		bt.data = nil
+		return nil
+	}
+
+	t, ok := pystate.(pickle.Tuple)
+	if !ok {
+		return fmt.Errorf("top: expect (...); got %T", pystate)
+	}
+	if !(1 <= len(t) && len(t) <= 2) {
+		return fmt.Errorf("top: expect [1..2](); got [%d]()", len(t))
+	}
+
+	// btree with 1 child bucket without oid
+	if len(t) == 1 {
+		t, ok := t[0].(pickle.Tuple)
+		if !ok {
+			return fmt.Errorf("bucket1: expect [1](); got %T", t[0])
+		}
+		if len(t) != 1 {
+			return fmt.Errorf("bucket1: expect [1](); got [%d]()", len(t))
+		}
+		bucket := zodb.NewPersistent(reflect.TypeOf(Bucket{}), bt.PJar()).(*Bucket)
+		err := (*bucketState)(bucket).PySetState(t[0])
+		if err != nil {
+			return fmt.Errorf("bucket1: %s", err)
+		}
+
+		bt.firstbucket = bucket
+		bt.data = []_BTreeItem{{key: 0, child: bucket}}
+		return nil
+	}
+
+	// regular btree
+	bt.firstbucket, ok = t[1].(*Bucket)
+	if !ok {
+		return fmt.Errorf("first bucket: must be Bucket; got %T", t[1])
+	}
+
+	t, ok = t[0].(pickle.Tuple)
+	if !ok {
+		return fmt.Errorf("data: expect (...); got %T", t[0])
+	}
+	if len(t)%2 == 0 {
+		return fmt.Errorf("data: expect [!%%2](); got [%d]()", len(t))
+	}
+
+	n := (len(t) + 1) / 2
+	bt.data = make([]_BTreeItem, 0, n)
+	for i, idx := 0, 0; i < n; i++ {
+		key := int64(0)
+		if i > 0 {
+			// key[0] is unused and not saved
+			key, ok = t[idx].(int64) // XXX Xint
+			if !ok {
+				return fmt.Errorf("data: [%d]: key must be integer; got %T", i, t[idx])
+			}
+			idx++
+		}
+		child := t[idx]
+		idx++
+
+		switch child.(type) {
+		default:
+			return fmt.Errorf("data: [%d]: child must be BTree|Bucket; got %T", i, child)
+
+		case *BTree:  // ok
+		case *Bucket: // ok
+		}
+
+		bt.data = append(bt.data, _BTreeItem{key: KEY(key), child: child})
+	}
+
+	return nil
+}
+
+
+// ---- register classes to ZODB ----
+
+func init() {
+	t := reflect.TypeOf
+	zodb.RegisterClass("BTrees.LOBTree.LOBucket", t(Bucket{}), t(bucketState{}))
+	zodb.RegisterClass("BTrees.LOBTree.LOBTree",  t(BTree{}),  t(btreeState{}))
+}

go/zodb/btree/btree_test.go

+// Copyright (C) 2018  Nexedi SA and Contributors.
+//                     Kirill Smelkov <kirr@nexedi.com>
+//
+// This program is free software: you can Use, Study, Modify and Redistribute
+// it under the terms of the GNU General Public License version 3, or (at your
+// option) any later version, as published by the Free Software Foundation.
+//
+// You can also Link and Combine this program with other software covered by
+// the terms of any of the Free Software licenses or any of the Open Source
+// Initiative approved licenses and Convey the resulting work. Corresponding
+// source of such a combination shall include the source code for all other
+// software used.
+//
+// This program is distributed WITHOUT ANY WARRANTY; without even the implied
+// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+//
+// See COPYING file for full licensing terms.
+// See https://www.nexedi.com/licensing for rationale and options.
+
+package btree
+
+//go:generate ./testdata/gen-testdata
+
+import (
+	"context"
+	"testing"
+
+	"lab.nexedi.com/kirr/go123/exc"
+	"lab.nexedi.com/kirr/neo/go/transaction"
+	"lab.nexedi.com/kirr/neo/go/zodb"
+	_ "lab.nexedi.com/kirr/neo/go/zodb/wks"
+)
+
+// kv is one (key, value) pair.
+type kv struct {
+	key   KEY
+	value interface{}
+}
+
+type bkind int
+const (
+	kindBucket bkind = iota
+	kindBTree
+)
+
+// testEntry is information about a Bucket or a BTree.
+type testEntry struct {
+	oid   zodb.Oid
+	kind  bkind
+	itemv []kv
+}
+
+// bmapping represents Get of Bucket or BTree.
+type bmapping interface {
+	Get(context.Context, KEY) (interface{}, bool, error)
+}
+
+func TestBTree(t *testing.T) {
+	X := exc.Raiseif
+	ctx := context.Background()
+	stor, err := zodb.OpenStorage(ctx, "testdata/1.fs", &zodb.OpenOptions{ReadOnly: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	db := zodb.NewDB(stor)
+
+	txn, ctx := transaction.New(ctx)
+	defer txn.Abort()
+
+	conn, err := db.Open(ctx, &zodb.ConnOptions{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// XXX close db/stor
+
+	// go through small test Buckets/BTrees and verify that Get(key) is as expected.
+	for _, tt := range smallTestv {
+		xobj, err := conn.Get(ctx, tt.oid)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		obj, ok := xobj.(bmapping)
+		if !ok {
+			t.Fatalf("%s: got %T;  want Bucket|BTree", tt.oid, xobj)
+		}
+
+		want := ""
+		switch tt.kind {
+		case kindBucket:
+			if _, ok = obj.(*Bucket); !ok {
+				want = "Bucket"
+			}
+		case kindBTree:
+			if _, ok = obj.(*BTree); !ok {
+				want = "BTree"
+			}
+		default:
+			panic(0)
+		}
+
+		if want != "" {
+			t.Fatalf("%s: got %T;  want %s", tt.oid, obj, want)
+		}
+
+		for _, kv := range tt.itemv {
+			value, ok, err := obj.Get(ctx, kv.key)
+			if err != nil {
+				t.Error(err)
+				continue
+			}
+
+			if !ok {
+				t.Errorf("%s: get %v -> ø;  want %v", tt.oid, kv.key, kv.value)
+				continue
+			}
+
+			if value != kv.value {
+				t.Errorf("%s: get %v -> %v;  want %v", tt.oid, kv.key, value, kv.value)
+			}
+
+			// XXX .next == nil
+			// XXX check keys, values directly (i.e. there is nothing else)
+		}
+	}
+
+
+	// B3 is a large BTree with {i: i} data.
+	// verify Get(key) and that different bucket links lead to the same in-RAM object.
+	xB3, err := conn.Get(ctx, B3_oid)
+	if err != nil {
+		t.Fatal(err)
+	}
+	B3, ok := xB3.(*BTree)
+	if !ok {
+		t.Fatalf("B3: %v; got %T;  want BTree", B3_oid, xB3)
+	}
+
+	for i := KEY(0); i <= B3_maxkey; i++ {
+		v, ok, err := B3.Get(ctx, i)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !ok {
+			t.Fatalf("B3: get %v -> ø;  want %v", i, i)
+		}
+		if int64(i) != v {
+			t.Fatalf("B3: get %v -> %v;  want %v", i, v, i)
+		}
+	}
+
+	// verifyFirstBucket verifies that b.firstbucket is correct and returns it.
+	var verifyFirstBucket func(b *BTree) *Bucket
+	verifyFirstBucket = func(b *BTree) *Bucket {
+		err := b.PActivate(ctx);	X(err)
+		defer b.PDeactivate()
+
+		var firstbucket *Bucket
+
+		switch child := b.data[0].child.(type) {
+		default:
+			t.Fatalf("btree(%s): child[0] is %T", b.POid(), b.data[0].child)
+
+		case *BTree:
+			firstbucket = verifyFirstBucket(child)
+
+		case *Bucket:
+			firstbucket = child
+		}
+
+		if firstbucket != b.firstbucket {
+			t.Fatalf("btree(%s): firstbucket -> %p (oid: %s); actual first bucket = %p (oid: %s)",
+				b.POid(), b.firstbucket, b.firstbucket.POid(), firstbucket, firstbucket.POid())
+		}
+
+		return firstbucket
+	}
+
+	verifyFirstBucket(B3)
+}

go/zodb/btree/testdata/.gitignore

+/*.lock
+/*.tmp
+/*.tr[0-9]

go/zodb/btree/testdata/gen-testdata

+#!/usr/bin/env python2
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018  Nexedi SA and Contributors.
+#                     Kirill Smelkov <kirr@nexedi.com>
+#
+# This program is free software: you can Use, Study, Modify and Redistribute
+# it under the terms of the GNU General Public License version 3, or (at your
+# option) any later version, as published by the Free Software Foundation.
+#
+# You can also Link and Combine this program with other software covered by
+# the terms of any of the Free Software licenses or any of the Open Source
+# Initiative approved licenses and Convey the resulting work. Corresponding
+# source of such a combination shall include the source code for all other
+# software used.
+#
+# This program is distributed WITHOUT ANY WARRANTY; without even the implied
+# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+#
+# See COPYING file for full licensing terms.
+# See https://www.nexedi.com/licensing for rationale and options.
+"""generate test data for btree serialization tests"""
+
+from ZODB.DB import DB
+from BTrees.LOBTree import LOBucket, LOBTree
+from ZODB.utils import u64
+import os, os.path, transaction
+
+from golang.gcompat import qq
+
+def rm_f(path):
+    if os.path.exists(path):
+        os.remove(path)
+
+def main():
+    import zodbtools.test.gen_testdata  # to make time predictable (XXX)
+
+    outfs = "testdata/1.fs"
+    rm_f(outfs)
+    rm_f(outfs + ".index")
+    db = DB(outfs)
+    conn = db.open()
+    root = conn.root()
+
+    root['b0'] = b0 = LOBucket()                         # empty bucket
+    root['b1'] = b1 = LOBucket([(10, 17)])               # 1k -> 1v
+    root['b2'] = b2 = LOBucket([(15, 1), (23, "hello")]) # 2k -> 2v
+
+    root['B0'] = B0 = LOBTree()                          # empty btree
+    root['B1'] = B1 = LOBTree({5: 4})                    # btree with 1 bucket (1kv)
+    root['B2'] = B2 = LOBTree({7: 3, 9: "world"})        # btree with 1 bucket (2kv)
+
+    root['B3'] = B3 = LOBTree(dict([(_, _) for _ in range(10000)]))
+
+    transaction.commit()
+
+    with open("ztestdata_expect_test.go", "w") as f:
+        def emit(v):
+            print >>f, v
+        emit("// Code generated by %s; DO NOT EDIT." % __file__)
+        emit("package btree\n")
+        #emit("import \"lab.nexedi.com/kirr/neo/go/zodb\"\n")
+
+        def emititems(b):
+            s = "testEntry{oid: %s, kind: %s, itemv: []kv{" \
+                    % (u64(b._p_oid), "kind%s" % type(b).__name__[2:])
+            for k, v in b.items():
+                if isinstance(v, str):
+                    v = qq(v)
+                elif isinstance(v, int):
+                    v = "int64(%d)" % v
+                else:
+                    raise RuntimeError("unsupported value type: %r" % v)
+
+                s += "{%s, %s}, " % (k, v)
+            s += "}},"
+            emit("\t"+s)
+
+        emit("\nvar smallTestv = [...]testEntry{")
+        for b in (b0, b1, b2, B0, B1, B2):
+            emititems(b)
+        emit("}")
+
+        emit("\nconst B3_oid = %s" % u64(B3._p_oid))
+        emit("const B3_maxkey = %d" % B3.maxKey())
+
+
+    conn.close()
+    db.close()
+
+
+if __name__ == '__main__':
+    main()

go/zodb/btree/ztestdata_expect_test.go

+// Code generated by ./testdata/gen-testdata; DO NOT EDIT.
+package btree
+
+
+var smallTestv = [...]testEntry{
+	testEntry{oid: 6, kind: kindBucket, itemv: []kv{}},
+	testEntry{oid: 3, kind: kindBucket, itemv: []kv{{10, int64(17)}, }},
+	testEntry{oid: 1, kind: kindBucket, itemv: []kv{{15, int64(1)}, {23, "hello"}, }},
+	testEntry{oid: 2, kind: kindBTree, itemv: []kv{}},
+	testEntry{oid: 7, kind: kindBTree, itemv: []kv{{5, int64(4)}, }},
+	testEntry{oid: 4, kind: kindBTree, itemv: []kv{{7, int64(3)}, {9, "world"}, }},
+}
+
+const B3_oid = 5
+const B3_maxkey = 9999

go/zodb/zodb.go

@@ -144,6 +144,10 @@
 // level types that are registered with state type providing PyStateful (see
 // RegisterClass) are automatically (de)serialized as Python pickles(*).
 //
+// An example of application-level type with ZODB/py compatibility can be seen in
+// package lab.nexedi.com/kirr/neo/go/zodb/btree which provides BTree handling
+// for ZODB/go.
+//
 // --------
 //
 // (*) for pickle support package github.com/kisielk/og-rek is used.