parse and cache pbf files

binary/serialize.go

@@ -107,3 +107,25 @@ func UnmarshalWay(data []byte) (way *element.Way, err error) {
 	way.TagsFromArray(pbfWay.Tags)
 	return way, nil
 }
+
+func MarshalRelation(relation *element.Relation) ([]byte, error) {
+	pbfRelation := &model.Relation{}
+	pbfRelation.Id = &relation.Id
+	//pbfRelation.Members = relation.Members
+	pbfRelation.Tags = relation.TagsAsArray()
+	return proto.Marshal(pbfRelation)
+}
+
+func UnmarshalRelation(data []byte) (relation *element.Relation, err error) {
+	pbfRelation := &model.Relation{}
+	err = proto.Unmarshal(data, pbfRelation)
+	if err != nil {
+		return nil, err
+	}
+
+	relation = &element.Relation{}
+	relation.Id = *pbfRelation.Id
+	//relation.Nodes = pbfRelation.Node
+	relation.TagsFromArray(pbfRelation.Tags)
+	return relation, nil
+}

cache/db.go

@@ -45,6 +45,10 @@ func (p *Cache) GetCoord(id int64) *element.Node {
 	if err != nil {
 		panic(err)
 	}
+	if data == nil {
+		return nil
+	}
+
 	node, err := binary.UnmarshalCoord(id, data)
 	if err != nil {
 		panic(err)
@@ -69,6 +73,9 @@ func (p *Cache) GetNode(id int64) *element.Node {
 	if err != nil {
 		panic(err)
 	}
+	if data == nil {
+		return nil
+	}
 	node, err := binary.UnmarshalNode(data)
 	if err != nil {
 		panic(err)
@@ -93,6 +100,9 @@ func (p *Cache) GetWay(id int64) *element.Way {
 	if err != nil {
 		panic(err)
 	}
+	if data == nil {
+		return nil
+	}
 	way, err := binary.UnmarshalWay(data)
 	if err != nil {
 		panic(err)
@@ -100,6 +110,33 @@ func (p *Cache) GetWay(id int64) *element.Way {
 	return way
 }
 
+func (p *Cache) PutRelation(relation *element.Relation) {
+	keyBuf := make([]byte, 8)
+	bin.PutVarint(keyBuf, int64(relation.Id))
+	data, err := binary.MarshalRelation(relation)
+	if err != nil {
+		panic(err)
+	}
+	p.db.Put(p.wo, keyBuf, data)
+}
+
+func (p *Cache) GetRelation(id int64) *element.Relation {
+	keyBuf := make([]byte, 8)
+	bin.PutVarint(keyBuf, int64(id))
+	data, err := p.db.Get(p.ro, keyBuf)
+	if err != nil {
+		panic(err)
+	}
+	if data == nil {
+		return nil
+	}
+	relation, err := binary.UnmarshalRelation(data)
+	if err != nil {
+		panic(err)
+	}
+	return relation
+}
+
 func (p *Cache) Close() {
 	p.db.Close()
 }

cache/db_test.go

@@ -44,19 +44,27 @@ func TestReadWriteNode(t *testing.T) {
 	defer os.RemoveAll(cache_dir)
 
 	cache := NewCache(cache_dir)
-	node := &element.Node{}
-	node.Id = 1
+	node := &element.Node{
+		OSMElem: element.OSMElem{
+			Id:   1234,
+			Tags: element.Tags{"foo": "bar"}},
+	}
 	cache.PutNode(node)
 	cache.Close()
 
 	cache = NewCache(cache_dir)
 	defer cache.Close()
 
-	data := cache.GetNode(1)
-
-	if data.Id != 1 {
-		t.Errorf("unexpected result of GetNode(1): %v", data)
+	data := cache.GetNode(1234)
+	if data.Id != 1234 || data.Tags["foo"] != "bar" {
+		t.Errorf("unexpected result of GetNode: %v", data)
 	}
+
+	data = cache.GetNode(99)
+	if data != nil {
+		t.Error("missing node not nil")
+	}
+
 }
 
 func TestReadWriteWay(t *testing.T) {
@@ -64,18 +72,41 @@ func TestReadWriteWay(t *testing.T) {
 	defer os.RemoveAll(cache_dir)
 
 	cache := NewCache(cache_dir)
-	way := &element.Way{}
-	way.Id = 1
+	way := &element.Way{
+		OSMElem: element.OSMElem{
+			Id:   1234,
+			Tags: element.Tags{"foo": "bar"}},
+		Nodes: []int64{942374923, 23948234},
+	}
 	cache.PutWay(way)
 	cache.Close()
 
 	cache = NewCache(cache_dir)
 	defer cache.Close()
 
-	data := cache.GetWay(1)
+	data := cache.GetWay(1234)
 
-	if data.Id != 1 {
-		t.Errorf("unexpected result of GetWay(1): %v", data)
+	if data.Id != 1234 || data.Tags["foo"] != "bar" {
+		t.Errorf("unexpected result of GetWay: %#v", data)
+	}
+	if len(data.Nodes) != 2 ||
+		data.Nodes[0] != 942374923 ||
+		data.Nodes[1] != 23948234 {
+		t.Errorf("unexpected result of GetWay: %#v", data)
+	}
+}
+
+func TestReadMissingWay(t *testing.T) {
+	cache_dir, _ := ioutil.TempDir("", "goposm_test")
+	defer os.RemoveAll(cache_dir)
+
+	cache := NewCache(cache_dir)
+	defer cache.Close()
+
+	data := cache.GetWay(1234)
+
+	if data != nil {
+		t.Errorf("missing way did not return nil")
 	}
 }
 
@@ -88,7 +119,10 @@ func BenchmarkWriteWay(b *testing.B) {
 	defer cache.Close()
 
 	b.StartTimer()
-	way := &element.Way{}
+	way := &element.Way{
+		OSMElem: element.OSMElem{Tags: element.Tags{"foo": "bar"}},
+		Nodes:   []int64{942374923, 23948234},
+	}
 	for i := 0; i < b.N; i++ {
 		way.Id = int64(i)
 		cache.PutWay(way)

parser.go

@@ -2,11 +2,84 @@ package main
 
 import (
 	"fmt"
+	"goposm/cache"
+	"goposm/element"
 	"goposm/parser"
 	"os"
+	"runtime"
+	"sync"
 )
 
+func parse(filename string) {
+	nodes := make(chan element.Node)
+	ways := make(chan element.Way)
+	relations := make(chan element.Relation)
+
+	positions := parser.PBFBlockPositions(filename)
+
+	waitParser := sync.WaitGroup{}
+	for i := 0; i < 4; i++ {
+		waitParser.Add(1)
+		go func() {
+			for pos := range positions {
+				parser.ParseBlock(pos, nodes, ways, relations)
+			}
+			waitParser.Done()
+		}()
+	}
+
+	waitCounter := sync.WaitGroup{}
+	waitCounter.Add(1)
+	go func() {
+		cache := cache.NewCache("/tmp/goposm/way.cache")
+		defer cache.Close()
+
+		wayCounter := 0
+		for way := range ways {
+			cache.PutWay(&way)
+			wayCounter += 1
+		}
+		fmt.Println("ways", wayCounter)
+		waitCounter.Done()
+	}()
+	waitCounter.Add(1)
+	go func() {
+		cache := cache.NewCache("/tmp/goposm/relation.cache")
+		defer cache.Close()
+
+		relationCounter := 0
+		for rel := range relations {
+			cache.PutRelation(&rel)
+			relationCounter += 1
+		}
+		fmt.Println("relations", relationCounter)
+		waitCounter.Done()
+	}()
+
+	cache := cache.NewCache("/tmp/goposm/node.cache")
+	defer cache.Close()
+	for i := 0; i < 4; i++ {
+		waitCounter.Add(1)
+		go func() {
+			nodeCounter := 0
+			for node := range nodes {
+				cache.PutNode(&node)
+				nodeCounter += 1
+			}
+			fmt.Println("nodes", nodeCounter)
+			waitCounter.Done()
+		}()
+	}
+	waitParser.Wait()
+	close(nodes)
+	close(ways)
+	close(relations)
+	waitCounter.Wait()
+}
+
 func main() {
-	parser.PBFStats(os.Args[1])
+	runtime.GOMAXPROCS(runtime.NumCPU())
+	parse(os.Args[1])
+	//parser.PBFStats(os.Args[1])
 	fmt.Println("done")
 }

parser/lowlevel.go

@@ -52,6 +52,7 @@ func (pbf *PBF) BlockPositions() (positions chan BlockPosition) {
 			offset, size := pbf.NextDataPosition()
 			if size == 0 {
 				close(positions)
+				pbf.Close()
 				return
 			}
 			positions <- BlockPosition{pbf.filename, offset, size}

parser/pbf.go

@@ -208,11 +208,38 @@ func PBFBlockPositions(filename string) chan BlockPosition {
 	if err != nil {
 		log.Fatal(err)
 	}
-	defer pbf.Close()
 
 	return pbf.BlockPositions()
 }
 
+func ParseBlock(pos BlockPosition, nodes chan element.Node, ways chan element.Way, relations chan element.Relation) {
+	block := ReadPrimitiveBlock(pos)
+	stringtable := NewStringTable(block.GetStringtable())
+
+	for _, group := range block.Primitivegroup {
+		dense := group.GetDense()
+		if dense != nil {
+			parsedNodes := ReadDenseNodes(dense, block, stringtable)
+			for _, node := range parsedNodes {
+				nodes <- node
+			}
+		}
+		parsedNodes := ReadNodes(group.Nodes, block, stringtable)
+		for _, node := range parsedNodes {
+			nodes <- node
+		}
+		parsedWays := ReadWays(group.Ways, block, stringtable)
+		for _, way := range parsedWays {
+			ways <- way
+		}
+		parsedRelations := ReadRelations(group.Relations, block, stringtable)
+		for _, rel := range parsedRelations {
+			relations <- rel
+		}
+	}
+
+}
+
 func PBFStats(filename string) {
 	pbf, err := Open(filename)
 	if err != nil {