Enable Dynamic Reader Offset Support for Specific Extraction of Nodes, Ways, and Edges

decode.go

@@ -18,6 +18,17 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
+type CountingReader struct {
+	r   io.Reader
+	Pos int64
+}
+
+func (cr *CountingReader) Read(p []byte) (int, error) {
+	n, err := cr.r.Read(p)
+	cr.Pos += int64(n)
+	return n, err
+}
+
 const (
 	maxBlobHeaderSize = 64 * 1024
 
@@ -102,9 +113,28 @@ type pair struct {
 	e error
 }
 
+type resultWithOffset struct {
+	pair   pair
+	offset int64
+}
+
+type PbfIndex struct {
+	FirstNodeOffset     int64
+	FirstWayOffset      int64
+	FirstRelationOffset int64
+	Nodes               int64
+	Ways                int64
+	Relations           int64
+}
+
+type workerState struct {
+	PbfIndex
+	processor OsmRawProcessor
+}
+
 // A Decoder reads and decodes OpenStreetMap PBF data from an input stream.
 type Decoder struct {
-	r          io.Reader
+	r          *CountingReader
 	serializer chan pair
 
 	buf *bytes.Buffer
@@ -115,20 +145,37 @@ type Decoder struct {
 	headerOnce sync.Once
 
 	// for data decoders
-	inputs  []chan<- pair
+	inputs  []chan<- resultWithOffset
 	outputs []<-chan pair
+
+	validateHeader bool
+	workerStates   []*workerState
+	processors     []OsmRawProcessor
+	wg             sync.WaitGroup
 }
 
 // NewDecoder returns a new decoder that reads from r.
 func NewDecoder(r io.Reader) *Decoder {
+	cr := &CountingReader{r: r}
 	d := &Decoder{
-		r:          r,
-		serializer: make(chan pair, 8000), // typical PrimitiveBlock contains 8k OSM entities
+		r:              cr,
+		serializer:     make(chan pair, 8000), // typical PrimitiveBlock contains 8k OSM entities
+		validateHeader: true,
 	}
 	d.SetBufferSize(initialBlobBufSize)
 	return d
 }
 
+func (dec *Decoder) WithValidateHeader(b bool) *Decoder {
+	dec.validateHeader = b
+	return dec
+}
+
+func (dec *Decoder) WithProcessors(processors []OsmRawProcessor) *Decoder {
+	dec.processors = processors
+	return dec
+}
+
 // SetBufferSize sets initial size of decoding buffer. Default value is 1MB, you can set higher value
 // (for example, MaxBlobSize) for (probably) faster decoding, or lower value for reduced memory consumption.
 // Any value will produce valid results; buffer will grow automatically if required.
@@ -142,56 +189,105 @@ func (dec *Decoder) Header() (*Header, error) {
 	return dec.header, dec.readOSMHeader()
 }
 
+type OsmRawProcessor func(objects []interface{}, err error)
+
 // Start decoding process using n goroutines.
 func (dec *Decoder) Start(n int) error {
-	if n < 1 {
+
+	useProcessors := false
+	if len(dec.processors) > 0 {
+		useProcessors = true
+		n = len(dec.processors)
+	} else if n < 1 {
 		n = 1
 	}
 
-	if err := dec.readOSMHeader(); err != nil {
-		return err
+	if dec.validateHeader {
+		if err := dec.readOSMHeader(); err != nil {
+			return err
+		}
 	}
 
+	dec.workerStates = make([]*workerState, 0, n)
+
 	// start data decoders
 	for i := 0; i < n; i++ {
-		input := make(chan pair)
+		input := make(chan resultWithOffset)
 		output := make(chan pair)
+
+		workerState := &workerState{
+			processor: dec.processors[i],
+		}
+		dec.workerStates = append(dec.workerStates, workerState)
+		dec.wg.Add(1)
 		go func() {
 			dd := new(dataDecoder)
 			for p := range input {
-				if p.e == nil {
+				if p.pair.e == nil {
 					// send decoded objects or decoding error
-					objects, err := dd.Decode(p.i.(*OSMPBF.Blob))
-					output <- pair{objects, err}
+					objects, err := dd.Decode(p.pair.i.(*OSMPBF.Blob))
+					for _, o := range objects {
+						switch o.(type) {
+						case *Node:
+							workerState.Nodes++
+							if workerState.FirstNodeOffset == 0 || p.offset < workerState.FirstNodeOffset {
+								workerState.FirstNodeOffset = p.offset
+							}
+						case *Way:
+							workerState.Ways++
+							if workerState.FirstWayOffset == 0 || p.offset < workerState.FirstWayOffset {
+								workerState.FirstWayOffset = p.offset
+							}
+						case *Relation:
+							workerState.Relations++
+							if workerState.FirstRelationOffset == 0 || p.offset < workerState.FirstRelationOffset {
+								workerState.FirstRelationOffset = p.offset
+							}
+						}
+					}
+					if useProcessors {
+						workerState.processor(objects, err)
+					} else {
+						output <- pair{objects, err}
+					}
 				} else {
 					// send input error as is
-					output <- pair{nil, p.e}
+					if useProcessors {
+						workerState.processor(nil, p.pair.e)
+					} else {
+						output <- pair{nil, p.pair.e}
+					}
 				}
 			}
+			dec.wg.Done()
 			close(output)
 		}()
 
 		dec.inputs = append(dec.inputs, input)
 		dec.outputs = append(dec.outputs, output)
 	}
 
+	if !useProcessors {
+		return nil
+	}
+
 	// start reading OSMData
 	go func() {
 		var inputIndex int
 		for {
 			input := dec.inputs[inputIndex]
 			inputIndex = (inputIndex + 1) % n
-
+			pos := dec.r.Pos
 			blobHeader, blob, err := dec.readFileBlock()
 			if err == nil && blobHeader.GetType() != "OSMData" {
 				err = fmt.Errorf("unexpected fileblock of type %s", blobHeader.GetType())
 			}
 			if err == nil {
 				// send blob for decoding
-				input <- pair{blob, nil}
+				input <- resultWithOffset{pair{blob, nil}, pos}
 			} else {
 				// send input error as is
-				input <- pair{nil, err}
+				input <- resultWithOffset{pair{nil, err}, pos}
 				for _, input := range dec.inputs {
 					close(input)
 				}
@@ -394,3 +490,77 @@ func (dec *Decoder) decodeOSMHeader(blob *OSMPBF.Blob) error {
 
 	return nil
 }
+
+// read offsets and counts
+
+func (dec *Decoder) NodeCount() int64 {
+	out := int64(0)
+	for _, state := range dec.workerStates {
+		out += state.Nodes
+	}
+	return out
+}
+
+func (dec *Decoder) WayCount() int64 {
+	out := int64(0)
+	for _, state := range dec.workerStates {
+		out += state.Ways
+	}
+	return out
+}
+
+func (dec *Decoder) RelationCount() int64 {
+	out := int64(0)
+	for _, state := range dec.workerStates {
+		out += state.Relations
+	}
+	return out
+}
+
+func (dec *Decoder) FirstNodeOffset() int64 {
+	out := int64(0)
+	for _, state := range dec.workerStates {
+		if out == 0 || state.FirstNodeOffset < out {
+			out = state.FirstNodeOffset
+		}
+	}
+	return out
+}
+
+func (dec *Decoder) FirstWayOffset() int64 {
+	out := int64(0)
+	for _, state := range dec.workerStates {
+		if out == 0 || state.FirstWayOffset < out {
+			out = state.FirstWayOffset
+		}
+	}
+	return out
+}
+
+func (dec *Decoder) FirstRelationOffset() int64 {
+	out := int64(0)
+	for _, state := range dec.workerStates {
+		if out == 0 || state.FirstRelationOffset < out {
+			out = state.FirstRelationOffset
+		}
+	}
+	return out
+}
+
+func (dec *Decoder) PbfIndex() PbfIndex {
+	return PbfIndex{
+		Nodes:               dec.NodeCount(),
+		Ways:                dec.WayCount(),
+		Relations:           dec.RelationCount(),
+		FirstNodeOffset:     dec.FirstNodeOffset(),
+		FirstWayOffset:      dec.FirstWayOffset(),
+		FirstRelationOffset: dec.FirstRelationOffset(),
+	}
+}
+
+func (dec *Decoder) Wait() {
+	if len(dec.processors) == 0 {
+		panic("Can only wait if processors are provided")
+	}
+	dec.wg.Wait()
+}