osmpbf: reuse memory if filter is false

Author: paulmach

osmpbf/decode_data.go

@@ -147,6 +147,9 @@ func (dec *dataDecoder) scanPrimitiveBlock(data []byte) error {
 func (dec *dataDecoder) scanPrimitiveGroup(data []byte) error {
 	msg := protoscan.New(data)
 
+	way := &osm.Way{Visible: true}
+	relation := &osm.Relation{Visible: true}
+
 	for msg.Next() {
 		fn := msg.FieldNumber()
 		if fn == 1 {
@@ -173,11 +176,20 @@ func (dec *dataDecoder) scanPrimitiveGroup(data []byte) error {
 				return err
 			}
 
-			err = dec.scanWays(data)
+			way, err = dec.scanWays(data, way)
 			if err != nil {
 				return err
 			}
 
+			if dec.scanner.FilterWay == nil || dec.scanner.FilterWay(way) {
+				dec.q = append(dec.q, way)
+				way = &osm.Way{Visible: true}
+			} else {
+				tags := way.Tags
+				nodes := way.Nodes
+				*way = osm.Way{Visible: true, Nodes: nodes[:0], Tags: tags[:0]}
+			}
+
 			continue
 		}
 
@@ -187,11 +199,20 @@ func (dec *dataDecoder) scanPrimitiveGroup(data []byte) error {
 				return err
 			}
 
-			err = dec.scanRelations(data)
+			relation, err = dec.scanRelations(data, relation)
 			if err != nil {
 				return err
 			}
 
+			if dec.scanner.FilterRelation == nil || dec.scanner.FilterRelation(relation) {
+				dec.q = append(dec.q, relation)
+				relation = &osm.Relation{Visible: true}
+			} else {
+				tags := relation.Tags
+				members := relation.Members
+				*relation = osm.Relation{Visible: true, Members: members[:0], Tags: tags[:0]}
+			}
+
 			continue
 		}
 
@@ -465,28 +486,21 @@ func (dec *dataDecoder) extractDenseNodes() error {
 				n.Tags = append(n.Tags, osm.Tag{Key: st[k], Value: st[v]})
 			}
 		}
-		if dec.scanner.NodeFilter == nil || dec.scanner.NodeFilter(*n) {
+
+		if dec.scanner.FilterNode == nil || dec.scanner.FilterNode(n) {
+			dec.q = append(dec.q, n)
+			n = &osm.Node{Visible: true}
+		} else {
 			// skip unwanted nodes
-			index++
-		}
-	}
-	if len(nodes) != index {
-		/* copy to an array just large enough */
-		tmp := make([]osm.Node, index)
-		copy(tmp, nodes)
-		for i := range tmp {
-			dec.q = append(dec.q, &tmp[i])
-		}
-	} else {
-		for i := range nodes {
-			dec.q = append(dec.q, &nodes[i])
+			tags := n.Tags
+			*n = osm.Node{Visible: true, Tags: tags[:0]}
 		}
 	}
 
 	return nil
 }
 
-func (dec *dataDecoder) scanWays(data []byte) error {
+func (dec *dataDecoder) scanWays(data []byte, way *osm.Way) (*osm.Way, error) {
 	st := dec.primitiveBlock.GetStringtable().GetS()
 	granularity := int64(dec.primitiveBlock.GetGranularity())
 	dateGranularity := int64(dec.primitiveBlock.GetDateGranularity())
@@ -496,7 +510,10 @@ func (dec *dataDecoder) scanWays(data []byte) error {
 
 	msg := protoscan.New(data)
 
-	way := &osm.Way{Visible: true}
+	if way == nil {
+		way = &osm.Way{Visible: true}
+	}
+
 	var foundKeys, foundVals bool
 	for msg.Next() {
 		var i64 int64
@@ -515,7 +532,7 @@ func (dec *dataDecoder) scanWays(data []byte) error {
 		case 4: // info
 			d, err := msg.MessageData()
 			if err != nil {
-				return err
+				return nil, err
 			}
 
 			info := protoscan.New(d)
@@ -524,38 +541,38 @@ func (dec *dataDecoder) scanWays(data []byte) error {
 				case 1:
 					v, err := info.Int32()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					way.Version = int(v)
 				case 2:
 					v, err := info.Int64()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					millisec := time.Duration(v*dateGranularity) * time.Millisecond
 					way.Timestamp = time.Unix(0, millisec.Nanoseconds()).UTC()
 				case 3:
 					v, err := info.Int64()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					way.ChangesetID = osm.ChangesetID(v)
 				case 4:
 					v, err := info.Int32()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					way.UserID = osm.UserID(v)
 				case 5:
 					v, err := info.Uint32()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					way.User = st[v]
 				case 6:
 					v, err := info.Bool()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					way.Visible = v
 				default:
@@ -564,12 +581,12 @@ func (dec *dataDecoder) scanWays(data []byte) error {
 			}
 
 			if info.Err() != nil {
-				return info.Err()
+				return nil, info.Err()
 			}
 		case 8: // refs or nodes
 			dec.nodes, err = msg.Iterator(dec.nodes)
 			if err != nil {
-				return err
+				return nil, err
 			}
 
 			var prev, index int64
@@ -579,7 +596,7 @@ func (dec *dataDecoder) scanWays(data []byte) error {
 			for dec.nodes.HasNext() {
 				v, err := dec.nodes.Sint64()
 				if err != nil {
-					return err
+					return nil, err
 				}
 				prev = v + prev // delta encoding
 				way.Nodes[index].ID = osm.NodeID(prev)
@@ -588,7 +605,7 @@ func (dec *dataDecoder) scanWays(data []byte) error {
 		case 9: // lat
 			dec.wlats, err = msg.Iterator(dec.wlats)
 			if err != nil {
-				return err
+				return nil, err
 			}
 
 			var prev, index int64
@@ -598,7 +615,7 @@ func (dec *dataDecoder) scanWays(data []byte) error {
 			for dec.wlats.HasNext() {
 				v, err := dec.wlats.Sint64()
 				if err != nil {
-					return err
+					return nil, err
 				}
 				prev = v + prev // delta encoding
 				way.Nodes[index].Lat = 1e-9 * float64(latOffset+(granularity*prev))
@@ -607,7 +624,7 @@ func (dec *dataDecoder) scanWays(data []byte) error {
 		case 10: // lon
 			dec.wlons, err = msg.Iterator(dec.wlons)
 			if err != nil {
-				return err
+				return nil, err
 			}
 
 			var prev, index int64
@@ -617,7 +634,7 @@ func (dec *dataDecoder) scanWays(data []byte) error {
 			for dec.wlons.HasNext() {
 				v, err := dec.wlons.Sint64()
 				if err != nil {
-					return err
+					return nil, err
 				}
 				prev = v + prev // delta encoding
 				way.Nodes[index].Lon = 1e-9 * float64(lonOffset+(granularity*prev))
@@ -628,25 +645,23 @@ func (dec *dataDecoder) scanWays(data []byte) error {
 		}
 
 		if err != nil {
-			return err
+			return nil, err
 		}
 	}
 
 	if msg.Err() != nil {
-		return msg.Err()
+		return nil, msg.Err()
 	}
 
 	if foundKeys && foundVals {
 		var err error
 		way.Tags, err = scanTags(st, dec.keys, dec.vals)
 		if err != nil {
-			return err
+			return nil, err
 		}
 	}
-	if dec.scanner.WayFilter == nil || dec.scanner.WayFilter(*way) {
-		dec.q = append(dec.q, way)
-	}
-	return nil
+
+	return way, nil
 }
 
 // Make relation members from stringtable and three parallel arrays of IDs.
@@ -693,13 +708,16 @@ func extractMembers(
 	return members, nil
 }
 
-func (dec *dataDecoder) scanRelations(data []byte) error {
+func (dec *dataDecoder) scanRelations(data []byte, relation *osm.Relation) (*osm.Relation, error) {
 	st := dec.primitiveBlock.GetStringtable().GetS()
 	dateGranularity := int64(dec.primitiveBlock.GetDateGranularity())
 
 	msg := protoscan.New(data)
 
-	relation := &osm.Relation{Visible: true}
+	if relation == nil {
+		relation = &osm.Relation{Visible: true}
+	}
+
 	var foundKeys, foundVals, foundRoles, foundMemids, foundTypes bool
 	for msg.Next() {
 		var i64 int64
@@ -718,7 +736,7 @@ func (dec *dataDecoder) scanRelations(data []byte) error {
 		case 4: // info
 			d, err := msg.MessageData()
 			if err != nil {
-				return err
+				return nil, err
 			}
 
 			info := protoscan.New(d)
@@ -727,38 +745,38 @@ func (dec *dataDecoder) scanRelations(data []byte) error {
 				case 1:
 					v, err := info.Int32()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					relation.Version = int(v)
 				case 2:
 					v, err := info.Int64()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					millisec := time.Duration(v*dateGranularity) * time.Millisecond
 					relation.Timestamp = time.Unix(0, millisec.Nanoseconds()).UTC()
 				case 3:
 					v, err := info.Int64()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					relation.ChangesetID = osm.ChangesetID(v)
 				case 4:
 					v, err := info.Int32()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					relation.UserID = osm.UserID(v)
 				case 5:
 					v, err := info.Uint32()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					relation.User = st[v]
 				case 6:
 					v, err := info.Bool()
 					if err != nil {
-						return err
+						return nil, err
 					}
 					relation.Visible = v
 				default:
@@ -767,7 +785,7 @@ func (dec *dataDecoder) scanRelations(data []byte) error {
 			}
 
 			if info.Err() != nil {
-				return info.Err()
+				return nil, info.Err()
 			}
 		case 8: // refs or nodes
 			dec.roles, err = msg.Iterator(dec.roles)
@@ -783,20 +801,20 @@ func (dec *dataDecoder) scanRelations(data []byte) error {
 		}
 
 		if err != nil {
-			return err
+			return nil, err
 		}
 	}
 
 	if msg.Err() != nil {
-		return msg.Err()
+		return nil, msg.Err()
 	}
 
 	// possible for relation to not have tags
 	if foundKeys && foundVals {
 		var err error
 		relation.Tags, err = scanTags(st, dec.keys, dec.vals)
 		if err != nil {
-			return err
+			return nil, err
 		}
 	}
 
@@ -805,13 +823,11 @@ func (dec *dataDecoder) scanRelations(data []byte) error {
 		var err error
 		relation.Members, err = extractMembers(st, dec.roles, dec.memids, dec.types)
 		if err != nil {
-			return err
+			return nil, err
 		}
 	}
-	if dec.scanner.RelationFilter == nil || dec.scanner.RelationFilter(*relation) {
-		dec.q = append(dec.q, relation)
-	}
-	return nil
+
+	return relation, nil
 }
 
 func scanTags(stringTable []string, keys, vals *protoscan.Iterator) (osm.Tags, error) {

osmpbf/scanner.go

@@ -21,14 +21,17 @@ var _ osm.Scanner = &Scanner{}
 // https://golang.org/pkg/bufio/#Scanner
 type Scanner struct {
 	// Skip element types that are not needed. The data is skipped
-	// at the encoded protobuf level, but each block still
-	// needs to be decompressed.
-	SkipNodes      bool
-	SkipWays       bool
-	SkipRelations  bool
-	NodeFilter     func(osm.Node) bool
-	WayFilter      func(osm.Way) bool
-	RelationFilter func(osm.Relation) bool
+	// at the encoded protobuf level, but each block still needs to be decompressed.
+	SkipNodes     bool
+	SkipWays      bool
+	SkipRelations bool
+
+	// If the Filter function is false, the element well be skipped
+	// at the decoding level. The functions should be fast, they block the
+	// decoder, there are `procs` number of concurrent decoders.
+	FilterNode     func(*osm.Node) bool
+	FilterWay      func(*osm.Way) bool
+	FilterRelation func(*osm.Relation) bool
 
 	ctx    context.Context
 	closed bool